Use Tensor.split in _disunit_gradient

Author: ValerianRey

src/torchjd/autojac/_jac_to_grad.py

@@ -88,13 +88,8 @@ def _unite_jacobians(jacobians: list[Tensor]) -> Tensor:
 def _disunite_gradient(
     gradient_vector: Tensor, jacobians: list[Tensor], tensors: list[TensorWithJac]
 ) -> list[Tensor]:
-    gradients = list[Tensor]()
-    start = 0
-    for jacobian, t in zip(jacobians, tensors, strict=True):
-        end = start + jacobian[0].numel()
-        current_gradient_vector = gradient_vector[start:end]
-        gradients.append(current_gradient_vector.view(t.shape))
-        start = end
+    gradient_vectors = gradient_vector.split([t.numel() for t in tensors])
+    gradients = [g.view(t.shape) for g, t in zip(gradient_vectors, tensors)]
     return gradients