Fix an embarrasing brainfart, reference mnist has 2 hidden layers

tenstorrent · Jan 31, 2025 · 3581c14 · 3581c14
1 parent 28a6569
commit 3581c14
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/tests/jax/models/mnist/mlp/test_mnist_mlp.py b/tests/jax/models/mnist/mlp/test_mnist_mlp.py
@@ -67,7 +67,7 @@ def training_tester(request) -> MNISTMLPTester:
 @pytest.mark.parametrize(
     "inference_tester",
     [
-        (512,),
+        (512, 512),
         # (128,),
         # (128, 128),
         # (192, 128),