[FIX] Tensor and generator devices for sampling (#76)

curvlinops/fisher.py

@@ -10,11 +10,11 @@
 from torch import (
     Generator,
     Tensor,
+    as_tensor,
     autograd,
     multinomial,
     normal,
     softmax,
-    tensor,
     zeros_like,
 )
 from torch.nn import CrossEntropyLoss, MSELoss, Parameter
@@ -187,7 +187,7 @@ def _matmat(self, M: ndarray) -> ndarray:
         Returns:
             Matrix-multiplication result ``mat @ M``.
         """
-        if self._generator is None:
+        if self._generator is None or self._generator.device != self._device:
             self._generator = Generator(device=self._device)
         self._generator.manual_seed(self._seed)
 
@@ -271,8 +271,9 @@ def sample_grad_output(self, output: Tensor) -> Tensor:
         C = output.shape[1]
 
         if isinstance(self._loss_func, MSELoss):
-            std = tensor(
-                sqrt(0.5 / C) if self._loss_func.reduction == "mean" else sqrt(0.5)
+            std = as_tensor(
+                sqrt(0.5 / C) if self._loss_func.reduction == "mean" else sqrt(0.5),
+                device=output.device,
             )
             return 2 * normal(zeros_like(output), std, generator=self._generator)
 

curvlinops/kfac.py

@@ -339,7 +339,7 @@ def _compute_kfac(self):
             )
 
         # loop over data set, computing the Kronecker factors
-        if self._generator is None:
+        if self._generator is None or self._generator.device != self._device:
             self._generator = Generator(device=self._device)
         self._generator.manual_seed(self._seed)