Merge pull request #572 from aai-institute/fix/571-missing-move-preco…

…nditioner-cg

Overwrite `to` method of `CgInfluence`, add `to` method to precondito…

CHANGELOG.md

@@ -2,17 +2,19 @@
 
 ## Unreleased
 
-### Fixed
-
-- Fixed missing move of tensors to model device in `EkfacInfluence` 
-  implementation [PR #570](https://github.com/aai-institute/pyDVL/pull/570)
-
 ### Added
 
 - Add a device fixture for `pytest`, which depending on the availability and 
   user input (`pytest --with-cuda`) resolves to cuda device
   [PR #574](https://github.com/aai-institute/pyDVL/pull/574)
 
+### Fixed
+
+- Fixed missing move of tensors to model device in `EkfacInfluence` 
+  implementation [PR #570](https://github.com/aai-institute/pyDVL/pull/570)
+- Missing move to device of `preconditioner` in `CgInfluence` implementation
+  [PR #572](https://github.com/aai-institute/pyDVL/pull/572)
+
 ## 0.9.1 - Bug fixes, logging improvement
 
 ### Fixed

src/pydvl/influence/torch/influence_function_model.py

@@ -706,7 +706,9 @@ def mat_mat(x: torch.Tensor):
         R = (rhs - mat_mat(X)).T
         Z = R if self.pre_conditioner is None else self.pre_conditioner.solve(R)
         P, _, _ = torch.linalg.svd(Z, full_matrices=False)
-        active_indices = torch.as_tensor(list(range(X.shape[-1])), dtype=torch.long)
+        active_indices = torch.as_tensor(
+            list(range(X.shape[-1])), dtype=torch.long, device=self.model_device
+        )
 
         maxiter = self.maxiter if self.maxiter is not None else len(rhs) * 10
         y_norm = torch.linalg.norm(rhs, dim=1)
@@ -758,6 +760,11 @@ def mat_mat(x: torch.Tensor):
 
         return X.T
 
+    def to(self, device: torch.device):
+        if self.pre_conditioner is not None:
+            self.pre_conditioner = self.pre_conditioner.to(device)
+        return super().to(device)
+
 
 class LissaInfluence(TorchInfluenceFunctionModel):
     r"""

src/pydvl/influence/torch/pre_conditioner.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 from abc import ABC, abstractmethod
 from typing import Callable, Optional
 
@@ -70,6 +72,11 @@ def solve(self, rhs: torch.Tensor):
     def _solve(self, rhs: torch.Tensor):
         pass
 
+    @abstractmethod
+    def to(self, device: torch.device) -> PreConditioner:
+        """Implement this to move the (potentially fitted) preconditioner to a
+        specific device"""
+
 
 class JacobiPreConditioner(PreConditioner):
     r"""
@@ -141,6 +148,11 @@ def _solve(self, rhs: torch.Tensor):
 
         return rhs * inv_diag.unsqueeze(-1)
 
+    def to(self, device: torch.device) -> JacobiPreConditioner:
+        if self._diag is not None:
+            self._diag = self._diag.to(device)
+        return self
+
 
 class NystroemPreConditioner(PreConditioner):
     r"""
@@ -233,3 +245,8 @@ def _solve(self, rhs: torch.Tensor):
             result = result.squeeze()
 
         return result
+
+    def to(self, device: torch.device) -> NystroemPreConditioner:
+        if self._low_rank_approx is not None:
+            self._low_rank_approx = self._low_rank_approx.to(device)
+        return self