Fix elif

Author: PierreQuinton

src/torchjd/autogram/_gramian_computer.py

@@ -111,7 +111,7 @@ def _compute_gramian(self, dY1: Tensor, dY2: Tensor, X: Tensor) -> Tensor:
             # 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)
-        if dY1.ndim == 2:
+        elif 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