Avoid inplicit cuda call

pomonam · Mar 19, 2024 · a1169fb · a1169fb
1 parent bf0cd76
commit a1169fb
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/kronfluence/module/tracked_module.py b/kronfluence/module/tracked_module.py
@@ -353,8 +353,9 @@ def synchronize_covariance_matrices(self) -> None:
         if dist.is_initialized() and torch.cuda.is_available() and self._covariance_matrices_available():
             # Note that only the main process holds the aggregated covariance matrix.
             for covariance_factor_name in COVARIANCE_FACTOR_NAMES:
+                self._storage[covariance_factor_name] = self._storage[covariance_factor_name].cuda()
                 dist.reduce(
-                    tensor=self._storage[covariance_factor_name].cuda(),
+                    tensor=self._storage[covariance_factor_name],
                     op=dist.ReduceOp.SUM,
                     dst=0,
                 )