Fix device mismatch when moving TransformedEmulator to a new device

Use in calibration workflows

Author: radka-j

autoemulate/calibration/bayes.py

@@ -85,7 +85,7 @@ def __init__(
             calibration_params = list(parameter_range.keys())
         self.calibration_params = calibration_params
         self.emulator = emulator
-        self.emulator.device = self.device
+        self.emulator.to(self.device)
         self.output_names = list(observations.keys())
         self.logger, self.progress_bar = get_configured_logger(log_level)
         self.logger.info(

autoemulate/calibration/history_matching.py

@@ -364,7 +364,7 @@ def __init__(
             raise ValueError(msg)
 
         self.transformed_emulator_params = transformed_emulator_params or {}
-        self.emulator.device = self.device
+        self.emulator.to(self.device)
 
         # New data is simulated in `run()` and appended here
         # It can be used to refit the emulator

autoemulate/calibration/interval_excursion_set.py

@@ -131,7 +131,7 @@ def __init__(
         self.calibration_params = list(parameters_range.keys())
         self.d = len(self.parameters_range)
         self.emulator = emulator
-        self.emulator.device = self.device
+        self.emulator.to(self.device)
         # TODO: we might want to check that the len equals the number of tasks returned
         self.output_names = output_names
         self.logger, self.progress_bar = get_configured_logger(log_level)

autoemulate/emulators/base.py

@@ -8,8 +8,9 @@
 from torch.distributions import TransformedDistribution
 from torch.optim.lr_scheduler import ExponentialLR, LRScheduler
 
-from autoemulate.core.device import TorchDeviceMixin
+from autoemulate.core.device import TorchDeviceMixin, get_torch_device
 from autoemulate.core.types import (
+    DeviceLike,
     DistributionLike,
     GaussianLike,
     NumpyLike,
@@ -37,6 +38,30 @@ class Emulator(ABC, ValidationMixin, ConversionMixin, TorchDeviceMixin):
     y_transform: StandardizeTransform | None = None
     supports_uq: bool = False
 
+    def to(self, device: DeviceLike) -> "Emulator":  # type: ignore[override]
+        """
+        Move the emulator to the given device.
+
+        Subclasses may override this to move additional state (e.g. transforms,
+        cached tensors). The base implementation updates ``self.device`` and, for
+        emulators that are also ``nn.Module`` instances, delegates to
+        ``nn.Module.to()`` to move parameters and buffers.
+
+        Parameters
+        ----------
+        device: DeviceLike
+            The target device.
+
+        Returns
+        -------
+        Emulator
+            ``self``, for method chaining.
+        """
+        self.device = get_torch_device(device)
+        if isinstance(self, nn.Module):
+            nn.Module.to(self, self.device)
+        return self
+
     @abstractmethod
     def _fit(self, x: TensorLike, y: TensorLike): ...
 
@@ -528,7 +553,9 @@ def predict(self, x: TensorLike, with_grad: bool = False) -> GaussianLike:
         return pred
 
 
-class PyTorchBackend(nn.Module, Emulator):
+class PyTorchBackend(  # type: ignore[reportIncompatibleMethodOverride]
+    nn.Module, Emulator
+):
     """
     `PyTorchBackend` provides a backend for PyTorch models.
 

autoemulate/emulators/transformed/base.py

@@ -1,9 +1,10 @@
 import torch
 from linear_operator.operators import DiagLinearOperator
+from torch import nn
 from torch.distributions import ComposeTransform, Transform, TransformedDistribution
 from torch.func import jacrev
 
-from autoemulate.core.device import TorchDeviceMixin
+from autoemulate.core.device import TorchDeviceMixin, get_torch_device
 from autoemulate.core.types import (
     DeviceLike,
     DistributionLike,
@@ -203,6 +204,59 @@ def _fit_transforms(self, x: TensorLike, y: TensorLike):
             all(self._y_transforms_affine) if self._y_transforms_affine else False
         )
 
+    def to(self, device: DeviceLike) -> "TransformedEmulator":
+        """
+        Move the emulator and all its state to the given device.
+
+        Moves the underlying model, transforms, and cached tensors to ``device``.
+
+        Parameters
+        ----------
+        device: DeviceLike
+            The target device (e.g. ``"cpu"``, ``"mps"``, ``"cuda"``).
+
+        Returns
+        -------
+        TransformedEmulator
+            ``self``, for method chaining.
+        """
+        device = get_torch_device(device)
+        self.device = device
+
+        # Move the underlying model (Emulator.to handles nn.Module delegation)
+        self.model.to(device)
+        # Clear cached prediction strategies that hold stale device refs
+        if hasattr(self.model, "_clear_cache"):
+            self.model._clear_cache()  # type: ignore[attr-defined]
+
+        # Move the inner model's own transforms (e.g. StandardizeTransform)
+        for attr in ("x_transform", "y_transform"):
+            transform = getattr(self.model, attr, None)
+            if transform is not None:
+                self._move_transform_to_device(transform, device)
+
+        # Move transform state tensors
+        for transform in self.x_transforms + self.y_transforms:
+            self._move_transform_to_device(transform, device)
+
+        # Move cached Jacobian
+        if self._fixed_jacobian_y_inv is not None:
+            self._fixed_jacobian_y_inv = self._fixed_jacobian_y_inv.to(device)
+
+        return self
+
+    @staticmethod
+    def _move_transform_to_device(transform: Transform, device: torch.device) -> None:
+        """Move a transform's tensor attributes to the given device."""
+        if isinstance(transform, nn.Module):
+            transform.to(device)
+        if hasattr(transform, "device"):
+            object.__setattr__(transform, "device", device)
+        for attr_name in list(vars(transform)):
+            val = getattr(transform, attr_name)
+            if isinstance(val, torch.Tensor):
+                setattr(transform, attr_name, val.to(device))
+
     def refit(self, x: TensorLike, y: TensorLike, retrain_transforms: bool = False):
         """
         Refit the emulator with new data and optionally retrain transforms.