Merge pull request #487 from aai-institute/bugfix/486-to-device-methods

Fix implementation of 'to' methods of TorchInfluenceFunctionModel imp…

Author: schroedk

CHANGELOG.md

@@ -6,6 +6,8 @@
 
 - Bug in using `DaskInfluenceCalcualator` with `TorchnumpyConverter`
   for single dimensional arrays [PR #485](https://github.com/aai-institute/pyDVL/pull/485)
+- Fix implementations of `to` methods of `TorchInfluenceFunctionModel` implementations
+  [PR #487](https://github.com/aai-institute/pyDVL/pull/487)
 
 ## 0.8.0 - 🆕 New interfaces, scaling computation, bug fixes and improvements 🎁
 

src/pydvl/influence/torch/influence_function_model.py

@@ -304,6 +304,16 @@ def influences_from_factors(
     def _solve_hvp(self, rhs: torch.Tensor) -> torch.Tensor:
         pass
 
+    def to(self, device: torch.device):
+        self.model = self.model.to(device)
+        self._model_params = {
+            k: p.detach().to(device)
+            for k, p in self.model.named_parameters()
+            if p.requires_grad
+        }
+        self._model_device = device
+        return self
+
 
 class DirectInfluence(TorchInfluenceFunctionModel):
     r"""
@@ -402,15 +412,9 @@ def _solve_hvp(self, rhs: torch.Tensor) -> torch.Tensor:
         ).T
 
     def to(self, device: torch.device):
-        self.hessian = self.hessian.to(device)
-        self.model = self.model.to(device)
-        self._model_device = device
-        self._model_params = {
-            k: p.detach().to(device)
-            for k, p in self.model.named_parameters()
-            if p.requires_grad
-        }
-        return self
+        if self.is_fitted:
+            self.hessian = self.hessian.to(device)
+        return super().to(device)
 
 
 class CgInfluence(TorchInfluenceFunctionModel):
@@ -537,16 +541,6 @@ def reg_hvp(v: torch.Tensor):
             batch_cg[idx] = batch_result
         return batch_cg
 
-    def to(self, device: torch.device):
-        self.model = self.model.to(device)
-        self._model_params = {
-            k: p.detach().to(device)
-            for k, p in self.model.named_parameters()
-            if p.requires_grad
-        }
-        self._model_device = device
-        return self
-
     @staticmethod
     def _solve_cg(
         hvp: Callable[[torch.Tensor], torch.Tensor],
@@ -873,6 +867,6 @@ def _solve_hvp(self, rhs: torch.Tensor) -> torch.Tensor:
         return result.t()
 
     def to(self, device: torch.device):
-        return ArnoldiInfluence(
-            self.model.to(device), self.loss, self.low_rank_representation.to(device)
-        )
+        if self.is_fitted:
+            self.low_rank_representation = self.low_rank_representation.to(device)
+        return super().to(device)