diff --git a/test/test_tiny.py b/test/test_tiny.py index 0c18e6a0a8..72f274a6b8 100644 --- a/test/test_tiny.py +++ b/test/test_tiny.py @@ -134,8 +134,8 @@ class TestTiny(unittest.TestCase): def test_mnist_backward(self): # NOTE: we don't have the whole model here for speed layers = [ - nn.Conv2d(1, 32, 5), Tensor.relu, - nn.Conv2d(32, 32, 5), Tensor.relu] + nn.Conv2d(1, 8, 5), Tensor.relu, + nn.Conv2d(8, 8, 5), Tensor.relu] # replace random weights with ones # TODO: there's a bug here where it's tying two of the biases together. we need UNIQUE const @@ -144,7 +144,7 @@ class TestTiny(unittest.TestCase): # realize gradients for x in nn.state.get_parameters(layers): x.requires_grad_() - Tensor.empty(4, 1, 28, 28).sequential(layers).sum().backward() + Tensor.empty(4, 1, 14, 14).sequential(layers).sum().backward() Tensor.realize(*[x.grad for x in nn.state.get_parameters(layers) if x.grad is not None]) # *** image ***