Update test

Author: Kade Lin

tests/test_nn.py

@@ -38,6 +38,9 @@ def test_max(t: Tensor) -> None:
             expected_max = max([t[i, j, k] for k in range(t.shape[2])])
             assert_close(out[i, j, 0], expected_max)
 
+
+@given(tensors(shape=(2, 3, 4)))
+def test_max_backward(t: Tensor) -> None:
     t.requires_grad_(True)
     t.zero_grad_()
     out = minitorch.max(t, dim=2)