Change shape of inputs

Author: ValerianRey

tests/unit/autojac/_transform/test_grad.py

@@ -68,13 +68,14 @@ def test_empty_outputs():
     `Iterable`.
     """
 
-    a = torch.tensor(1.0, requires_grad=True)
+    a1 = torch.tensor(1.0, requires_grad=True)
+    a2 = torch.tensor([1.0, 2.0], requires_grad=True)
     input = Gradients({})
 
-    grad = Grad(outputs=[], inputs=[a])
+    grad = Grad(outputs=[], inputs=[a1, a2])
 
     gradients = grad(input)
-    expected_gradients = {a: torch.zeros_like(a)}
+    expected_gradients = {a1: torch.zeros_like(a1), a2: torch.zeros_like(a2)}
 
     assert_tensor_dicts_are_close(gradients, expected_gradients)
 

tests/unit/autojac/_transform/test_jac.py

@@ -80,7 +80,7 @@ def test_empty_outputs(chunk_size: int | None):
     """
 
     a1 = torch.tensor(1.0, requires_grad=True)
-    a2 = torch.tensor(1.0, requires_grad=True)
+    a2 = torch.tensor([1.0, 2.0], requires_grad=True)
     input = Jacobians({})
 
     jac = Jac(outputs=[], inputs=[a1, a2], chunk_size=chunk_size)