Merge branch 'feature/ekfac_new_framework' into feature/ekfac_notebook_new_framework

Author: Xuzzo

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)
 - Implement new method: `EkfacInfluence`
   [PR #451](https://github.com/aai-institute/pyDVL/issues/451)
 

src/pydvl/influence/torch/influence_function_model.py

@@ -309,6 +309,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"""
@@ -407,15 +417,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):
@@ -542,16 +546,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],
@@ -878,9 +872,9 @@ 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)
 
 
 class EkfacInfluence(TorchInfluenceFunctionModel):
@@ -1480,7 +1474,6 @@ def explore_hessian_regularization(
         return influences_by_reg_value
 
     def to(self, device: torch.device):
-        self.model.to(device)
         if self.is_fitted:
             self.ekfac_representation.to(device)
-        return self
+        return super().to(device)