ndim=1 cannot happen (we are batched for now).

Author: PierreQuinton

src/torchjd/autogram/_gramian_computer.py

@@ -107,11 +107,7 @@ def _compute_gramian(self, dY1: Tensor, dY2: Tensor, X: Tensor) -> Tensor:
         """
 
         # TODO: add support for ndim==4 or find solution that works for any ndim.
-        if dY1.ndim == 1:
-            # TODO: not sure that this even works
-            G_b = torch.einsum("k,k->", dY1, dY2)
-            G_W = torch.einsum("k,l,l,k->", dY1, X, X, dY2)
-        elif dY1.ndim == 2:
+        if dY1.ndim == 2:
             G_b = torch.einsum("ak,ik->ai", dY1, dY2)
             G_W = torch.einsum("ak,al,il,ik->ai", dY1, X, X, dY2)
         elif dY1.ndim == 3:  # Typical in transformers