Add missing type annotation for chunk_size

Author: ValerianRey

tests/unit/autojac/test_backward.py

@@ -268,7 +268,7 @@ def test_multiple_tensors() -> None:
 
 
 @mark.parametrize("chunk_size", [None, 1, 2, 4])
-def test_various_valid_chunk_sizes(chunk_size) -> None:
+def test_various_valid_chunk_sizes(chunk_size: int | None) -> None:
     """Tests that backward works for various valid values of parallel_chunk_size."""
 
     a1 = tensor_([1.0, 2.0], requires_grad=True)

tests/unit/autojac/test_jac.py

@@ -268,7 +268,7 @@ def test_multiple_tensors() -> None:
 
 
 @mark.parametrize("chunk_size", [None, 1, 2, 4])
-def test_various_valid_chunk_sizes(chunk_size) -> None:
+def test_various_valid_chunk_sizes(chunk_size: int | None) -> None:
     """Tests that jac works for various valid values of parallel_chunk_size."""
 
     a1 = tensor_([1.0, 2.0], requires_grad=True)

tests/unit/autojac/test_mtl_backward.py

@@ -378,7 +378,7 @@ def test_non_scalar_loss_fails() -> None:
 
 
 @mark.parametrize("chunk_size", [None, 1, 2, 4])
-def test_various_valid_chunk_sizes(chunk_size) -> None:
+def test_various_valid_chunk_sizes(chunk_size: int | None) -> None:
     """Tests that mtl_backward works for various valid values of parallel_chunk_size."""
 
     p0 = tensor_([1.0, 2.0], requires_grad=True)