Make dependence on gramian explicit in CAGrad

Author: PierreQuinton

src/torchjd/aggregation/cagrad.py

@@ -73,12 +73,15 @@ def __init__(self, c: float, norm_eps: float):
 
     def forward(self, matrix: Tensor) -> Tensor:
         gramian = normalize(compute_gramian(matrix), self.norm_eps)
+        return self._compute_from_gramian(gramian)
+
+    def _compute_from_gramian(self, gramian: Tensor) -> Tensor:
         U, S, _ = torch.svd(gramian)
 
         reduced_matrix = U @ S.sqrt().diag()
         reduced_array = reduced_matrix.cpu().detach().numpy().astype(np.float64)
 
-        dimension = matrix.shape[0]
+        dimension = gramian.shape[0]
         reduced_g_0 = reduced_array.T @ np.ones(dimension) / dimension
         sqrt_phi = self.c * np.linalg.norm(reduced_g_0, 2)
 
@@ -97,6 +100,6 @@ def forward(self, matrix: Tensor) -> Tensor:
             # We are approximately on the pareto front
             weight_array = np.zeros(dimension)
 
-        weights = torch.from_numpy(weight_array).to(device=matrix.device, dtype=matrix.dtype)
+        weights = torch.from_numpy(weight_array).to(device=gramian.device, dtype=gramian.dtype)
 
         return weights