deepinv · jscanvic · Mar 3, 2026 · Dec 5, 2025 · Dec 8, 2025 · Dec 12, 2025
diff --git a/deepinv/models/dynamic.py b/deepinv/models/dynamic.py
@@ -85,8 +85,15 @@ def forward(self, y, physics: TimeMixin, **kwargs):
         :param y: measurements
         :parameter physics: forward operator acting on dynamic inputs
         """
+
+        static_physics = (
+            physics.to_static(device=y.device)
+            if hasattr(physics, "to_static")
+            else physics
+        )
+
         return self.backbone_net(
             self.average(y, getattr(physics, "mask", None)),
-            getattr(physics, "to_static", lambda: physics)(),
+            static_physics,
             **kwargs,
         )
diff --git a/deepinv/models/ram.py b/deepinv/models/ram.py
@@ -332,7 +332,9 @@ def forward(
             y = y / rescale_val.view([y.shape[0]] + [1] * (y.ndim - 1))
 
         if physics is None:
-            physics = dinv.physics.Denoising(noise_model=dinv.physics.ZeroNoise())
+            physics = dinv.physics.Denoising(
+                noise_model=dinv.physics.ZeroNoise(), device=y.device
+            )
 
         if img_size is None:
             if hasattr(physics, "img_shape") and physics.img_shape is not None:
@@ -354,7 +356,7 @@ def forward(
 
         use_pad = False
         if pad[0] != 0 or pad[1] != 0:
-            physics = PhysicsCropper(physics, pad)
+            physics = PhysicsCropper(physics, pad, device=y.device)
             use_pad = True
 
         x_in = physics.A_adjoint(y)

diff --git a/deepinv/physics/blur.py b/deepinv/physics/blur.py
diff --git a/deepinv/physics/cassi.py b/deepinv/physics/cassi.py
@@ -70,7 +70,7 @@ def __init__(
         rng: torch.Generator = None,
         **kwargs,
     ):
-        super().__init__(**kwargs)
+        super().__init__(device=device, **kwargs)
 
         if len(img_size) != 3:
             raise ValueError("img_size must be (C, H, W)")

diff --git a/deepinv/physics/compressed_sensing.py b/deepinv/physics/compressed_sensing.py
@@ -110,21 +110,20 @@ def __init__(
         rng: torch.Generator = None,
         **kwargs,
     ):
-        super().__init__(**kwargs)
+        super().__init__(device=device, **kwargs)
         self.name = f"CS_m{m}"
         self.img_size = img_size
         self.fast = fast
         self.channelwise = channelwise
         self.dtype = dtype
-        self.device = device
 
         if rng is None:
             self.rng = torch.Generator(device=device)
         else:
             # Make sure that the random generator is on the same device as the physic generator
             assert rng.device == torch.device(
                 device
-            ), f"The random generator is not on the same device as the Physics Generator. Got random generator on {rng.device} and the Physics Generator on {self.device}."
+            ), f"The random generator is not on the same device as the Physics Generator. Got random generator on {rng.device} and the Physics Generator on {device}."
             self.rng = rng
         self.register_buffer("initial_random_state", self.rng.get_state())
 

diff --git a/deepinv/physics/forward.py b/deepinv/physics/forward.py
@@ -260,7 +260,20 @@ def update_parameters(self, **kwargs):
                     and hasattr(self, key)
                     and isinstance(value, torch.Tensor)
                 ):
-                    self.register_buffer(key, value)
+
+                    if isinstance(getattr(self, key), torch.Tensor):
+                        if value.device.type != getattr(self, key).device.type:
+                            warnings.warn(
+                                f"The provided tensor for parameter '{key}' is on a different device ({value.device}) than the current parameter device ({getattr(self, key).device}). The current device will be used.",
+                                stacklevel=2,
+                            )
+
+                    # Move `value` to the buffer's device before updating
+                    # regardless of where the `value` tensor is located.
+                    # Also performs type casting if necessary.
+                    # If getattr(self, key) is None, torch.Tensor.to will
+                    # ignore the call and just return the original tensor.
+                    setattr(self, key, value.to(getattr(self, key)))
 
     # NOTE: Physics instances can hold instances of torch.Generator as
     # (possibly nested) attributes and they cannot be copied using deepcopy
@@ -373,6 +386,7 @@ class LinearPhysics(Physics):
         is used for computing it, and this parameter fixes the relative tolerance of the least squares algorithm.
     :param str solver: least squares solver to use. Choose between `'CG'`, `'lsqr'`, `'BiCGStab'` and `'minres'`. See :func:`deepinv.optim.linear.least_squares` for more details.
     :param bool implicit_backward_solver: If `True`, uses implicit differentiation for computing gradients through the :meth:`deepinv.physics.LinearPhysics.A_dagger` and :meth:`deepinv.physics.LinearPhysics.prox_l2`, using :func:`deepinv.optim.linear.least_squares_implicit_backward` instead of :func:`deepinv.optim.linear.least_squares`. This can significantly reduce memory consumption, especially when using many iterations. If `False`, uses the standard autograd mechanism, which can be memory-intensive. Default is `True`.
+    :param torch.device, str device: cpu or cuda, every registered buffer and module parameters are recursively pushed onto the device during initialization.
 
     |sep|
 
@@ -436,6 +450,7 @@ def __init__(
         tol=1e-4,
         solver="lsqr",
         implicit_backward_solver: bool = True,
+        device: torch.device | str = "cpu",
         **kwargs,
     ):
         super().__init__(
@@ -457,6 +472,37 @@ def __init__(
                 "Using implicit_backward_solver with a low number of iterations may produce inaccurate gradients during the backward pass. If you are not doing backpropagation through `A_dagger` or `prox_l2`, ignore this message. If you are training unfolded models, consider increasing max_iter."
             )
 
+        device_holder = torch.tensor(0.0, device=device)
+        self.register_buffer("_device_holder", device_holder, persistent=False)
+        # pushes all parameters/buffers to the specified device, including `noise_model`
+        self.to(device)
+
+    @property
+    def device(self) -> torch.device | str:
+        r"""
+        Returns the device where the physics parameters/buffers are stored.
+
+        :return: device of the physics parameters.
+        """
+        warnings.warn(
+            "Following torch.nn.Module's design, the 'device' attribute is deprecated and will be removed in a future version. To move the module's buffers/parameters to a different device, use the `to()` method."
+        )
+
+        return self._device_holder.device
+
+    @device.setter
+    def device(self, value: torch.device | str):
+        r"""
+        Sets the device where the physics parameters/buffers are stored.
+
+        :param device: device to which the physics parameters will be moved.
+        """
+        warnings.warn(
+            "Following torch.nn.Module's design, the 'device' attribute is deprecated and will be removed in a future version, i.e. doing `physics.device = device` will no longer work and throw an `AttributeError`. Use `physics.to(device)` instead."
+        )
+
+        self.to(value)
+
     def A_adjoint(self, y, **kwargs):
         r"""
         Computes transpose of the forward operator :math:`\tilde{x} = A^{\top}y`.
@@ -964,7 +1010,8 @@ class DecomposablePhysics(LinearPhysics):
         from the `V_adjoint` function and the `img_size` parameter.
         This automatic adjoint is computed using automatic differentiation, which is slower than a closed form adjoint, and can
         have a larger memory footprint. If you want to use the automatic adjoint, you should set the `img_size` parameter.
-    :param torch.nn.parameter.Parameter, float params: Singular values of the transform
+    :param torch.nn.parameter.Parameter, float mask: Singular values of the transform
+    :param torch.device, str device: cpu or cuda, every registered buffer and module parameters are recursively pushed onto the device during initialization.
 
     |sep|
 
@@ -1002,9 +1049,10 @@ def __init__(
         U_adjoint=None,
         V=None,
         mask=1.0,
+        device: torch.device | str = "cpu",
         **kwargs,
     ):
-        super().__init__(**kwargs)
+        super().__init__(device=device, **kwargs)
 
         assert not (
             U is None and not (U_adjoint is None)
@@ -1027,6 +1075,8 @@ def __init__(
         self.img_size = img_size
         self.register_buffer("mask", mask)
 
+        self.to(device)
+
     def A(self, x, mask=None, **kwargs) -> Tensor:
         r"""
         Applies the forward operator :math:`y = A(x)`.
@@ -1210,6 +1260,7 @@ class Denoising(DecomposablePhysics):
     The linear operator is just the identity mapping :math:`A(x)=x`
 
     :param torch.nn.Module noise: noise distribution, e.g., :class:`deepinv.physics.GaussianNoise`, or a user-defined torch.nn.Module. By default, it is set to Gaussian noise with a standard deviation of 0.1.
+    :param torch.device, str device: cpu or cuda, every registered buffer and module parameters are recursively pushed onto the device during initialization.
 
     |sep|
 
@@ -1229,10 +1280,23 @@ class Denoising(DecomposablePhysics):
 
     """
 
-    def __init__(self, noise_model: NoiseModel | None = None, **kwargs):
+    def __init__(
+        self,
+        noise_model: NoiseModel | None = None,
+        device: str | torch.device = "cpu",
+        **kwargs,
+    ):
         if noise_model is None:
             noise_model = GaussianNoise(sigma=0.1)
-        super().__init__(noise_model=noise_model, **kwargs)
+
+        if noise_model.rng is not None:
+            if noise_model.rng.device != device:
+                warnings.warn(
+                    f"argument `device`={device} is different from the random generator device of the noise model, `noise_model.rng.device`={noise_model.rng.device}. This will likely lead to errors during execution. The device argument will be ignored in favor of `noise_model.rng.device`={noise_model.rng.device}."
+                )
+                device = noise_model.rng.device
+
+        super().__init__(noise_model=noise_model, device=device, **kwargs)
 
 
 def adjoint_function(A, input_size, device="cpu", dtype=torch.float):

diff --git a/deepinv/physics/functional/radon.py b/deepinv/physics/functional/radon.py
@@ -72,9 +72,8 @@ def deg2rad(x: int | float | torch.Tensor) -> torch.Tensor:
 
 
 class AbstractFilter(nn.Module):
-    def __init__(self, device="cpu", dtype=torch.float):
+    def __init__(self, dtype: torch.dtype = torch.float32):
         super().__init__()
-        self.device = device
         self.dtype = dtype
 
     def forward(self, x: torch.Tensor, dim: int = -2) -> torch.Tensor:
@@ -98,7 +97,7 @@ def forward(self, x: torch.Tensor, dim: int = -2) -> torch.Tensor:
         )
         pad_width = projection_size_padded - input_size
 
-        f = self._get_fourier_filter(projection_size_padded).to(x.device)
+        f = self._get_fourier_filter(projection_size_padded, device=x.device)
         fourier_filter = self.create_filter(f)
         if dim == 2 or dim == -2:
             fourier_filter = fourier_filter.unsqueeze(-1)
@@ -149,12 +148,12 @@ def filter(
         elif dim == 3 or dim == -1:
             return result[:, :, :, :input_size]
 
-    def _get_fourier_filter(self, size):
+    def _get_fourier_filter(self, size, device: str | torch.device = "cpu"):
         n = torch.cat(
             [torch.arange(1, size / 2 + 1, 2), torch.arange(size / 2 - 1, 0, -2)]
         )
 
-        f = torch.zeros(size, dtype=self.dtype, device=self.device)
+        f = torch.zeros(size, dtype=self.dtype, device=device)
         f[0] = 0.25
         f[1::2] = -1 / (torch.pi * n) ** 2
 
@@ -216,7 +215,7 @@ def __init__(
     ):
         super().__init__()
         self.circle = circle
-        theta = theta if theta is not None else torch.arange(180).to(self.device)
+        theta = theta if theta is not None else torch.arange(180, device=device)
         self.register_buffer("theta", theta, persistent=False)
         self.dtype = dtype
         self.parallel_computation = parallel_computation
@@ -239,7 +238,7 @@ def __init__(
             if not "detector_spacing" in self.fan_parameters.keys():
                 self.fan_parameters["detector_spacing"] = 0.077
 
-        all_grids = self._create_grids(self.theta, in_size, circle).to(device)
+        all_grids = self._create_grids(self.theta, in_size, circle, device=device)
         if self.parallel_computation:
             self.register_buffer(
                 "all_grids",
@@ -308,7 +307,11 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
         return out
 
     def _create_grids(
-        self, angles: torch.Tensor, grid_size: int, circle: bool, device: str = "cpu"
+        self,
+        angles: torch.Tensor,
+        grid_size: int,
+        circle: bool,
+        device: str | torch.device = "cpu",
     ) -> torch.Tensor:
         if not circle:
             grid_size = int((SQRT2 * grid_size).ceil())
@@ -365,19 +368,18 @@ def __init__(
     ):
         super().__init__()
         self.circle = circle
-        self.device = device
-        theta = theta if theta is not None else torch.arange(180).to(self.device)
+        theta = theta if theta is not None else torch.arange(180, device=device)
         self.register_buffer("theta", theta, persistent=False)
         self.out_size = out_size
         self.in_size = in_size
         self.parallel_computation = parallel_computation
         self.dtype = dtype
 
-        ygrid, xgrid = self._create_yxgrid(in_size, circle)
+        ygrid, xgrid = self._create_yxgrid(in_size, circle, device=device)
         self.register_buffer("xgrid", xgrid, persistent=False)
         self.register_buffer("ygrid", ygrid, persistent=False)
 
-        all_grids = self._create_grids(self.theta, in_size, circle).to(self.device)
+        all_grids = self._create_grids(self.theta, in_size, circle, device=device)
         if self.parallel_computation:
             self.register_buffer(
                 "all_grids",
@@ -387,11 +389,7 @@ def __init__(
         else:
             self.register_buffer("all_grids", all_grids, persistent=False)
 
-        self.filter = (
-            RampFilter(dtype=self.dtype, device=self.device)
-            if use_filter
-            else lambda x: x
-        )
+        self.filter = RampFilter(dtype=self.dtype) if use_filter else lambda x: x
 
     def forward(self, x: torch.Tensor, filtering: bool = True) -> torch.Tensor:
         r"""
@@ -414,7 +412,7 @@ def forward(self, x: torch.Tensor, filtering: bool = True) -> torch.Tensor:
                 ch_size,
                 it_size,
                 it_size,
-                device=self.device,
+                device=x.device,
                 dtype=self.dtype,
             )
             for i_theta in range(len(self.theta)):
@@ -450,11 +448,11 @@ def forward(self, x: torch.Tensor, filtering: bool = True) -> torch.Tensor:
         return reco
 
     def _create_yxgrid(
-        self, in_size: int, circle: bool
+        self, in_size: int, circle: bool, device: str | torch.device = "cpu"
     ) -> tuple[torch.Tensor, torch.Tensor]:
         if not circle:
             in_size = int((SQRT2 * in_size).ceil())
-        unitrange = torch.linspace(-1, 1, in_size, dtype=self.dtype, device=self.device)
+        unitrange = torch.linspace(-1, 1, in_size, dtype=self.dtype, device=device)
         ygrid, xgrid = torch.meshgrid(unitrange, unitrange, indexing="ij")
         return ygrid, xgrid
 
@@ -463,14 +461,18 @@ def _XYtoT(self, theta: int | float | torch.Tensor) -> torch.Tensor:
         return T
 
     def _create_grids(
-        self, angles: torch.Tensor, grid_size: int, circle: bool
+        self,
+        angles: torch.Tensor,
+        grid_size: int,
+        circle: bool,
+        device: str | torch.device = "cpu",
     ) -> torch.Tensor:
         if not circle:
             grid_size = int((SQRT2 * grid_size).ceil())
         all_grids = []
         for i_theta in range(len(angles)):
             X = (
-                torch.ones(grid_size, dtype=self.dtype, device=self.device)
+                torch.ones(grid_size, dtype=self.dtype, device=device)
                 .view(-1, 1)
                 .repeat(1, grid_size)
                 * i_theta

diff --git a/deepinv/physics/generator/base.py b/deepinv/physics/generator/base.py
@@ -64,14 +64,13 @@ def __init__(
         self.step_func = step
         self.kwargs = kwargs
         self.factory_kwargs = {"device": device, "dtype": dtype}
-        self.device = device
         if rng is None:
             self.rng = torch.Generator(device=device)
         else:
             # Make sure that the random generator is on the same device as the physics generator
             assert rng.device == torch.device(
                 device
-            ), f"The random generator is not on the same device as the Physics Generator. Got random generator on {rng.device} and the Physics Generator named {self.__class__.__name__} on {self.device}."
+            ), f"The random generator is not on the same device as the Physics Generator. Got random generator on {rng.device} and the Physics Generator named {self.__class__.__name__} on {device}."
             self.rng = rng
 
         # NOTE: There is no use in moving RNG states from one device to another
@@ -88,6 +87,10 @@ def __init__(
         for k, v in kwargs.items():
             setattr(self, k, v)
 
+    @property
+    def device(self) -> torch.device:
+        return self.rng.device
+
     def step(self, batch_size: int = 1, seed: int = None, **kwargs) -> dict:
         r"""
         Generates a batch of parameters for the forward operator.

diff --git a/deepinv/physics/inpainting.py b/deepinv/physics/inpainting.py
@@ -96,7 +96,7 @@ def __init__(
         rng: torch.Generator = None,
         **kwargs,
     ):
-        super().__init__(**kwargs)
+        super().__init__(device=device, **kwargs)
 
         if isinstance(mask, torch.Tensor):
             mask = mask.to(device)