Skip to content

Wrong device in R3Refinement callback #744

@GiovanniCanali

Description

@GiovanniCanali

Describe the bug
Current points and newly generated samples are always stored on the CPU, causing issues whenever the chosen accelerator is different.

RuntimeError: Expected all tensors to be on the same device

Additionally, the docstring is not aligned with PINA conventions, and some variable names are not informative, making debugging more difficult.

To Reproduce

from pina.problem.zoo import Poisson2DSquareProblem
from pina.callback import R3Refinement
from pina.model import FeedForward
from pina.solver import PINN
from pina import Trainer

# Problem
problem = Poisson2DSquareProblem()
problem.discretise_domain(n=10)

# Model and Solver
model = FeedForward(2, 1, 20, 2)
solver = PINN(problem, model)

# Callback and Trainer
callback = R3Refinement(sample_every=10)
trainer = Trainer(solver, callbacks=callback, accelerator="gpu", max_epochs=25)
trainer.train()

Expected behavior
All tensors should consistently be moved to and kept on the same device as the selected accelerator.

Metadata

Metadata

Labels

bugSomething isn't working

Type

Projects

No projects

Milestone

No milestone

Relationships

None yet

Development

No branches or pull requests

Issue actions