fix tests.

0x00-pl · Nov 29, 2024 · 286b3f6 · 286b3f6
1 parent d3e4753
commit 286b3f6
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/tests/module_pool/simple_nn.py b/tests/module_pool/simple_nn.py
@@ -16,7 +16,7 @@ def forward(self, x):
 
 
 def check_torch_compile_forward(model, compiled_model):
-    input_data = torch.randn(2, 5)
+    input_data = torch.randn(1, 10)
     expected_output = model(input_data)
     actual_output = compiled_model(input_data)
     assert torch.allclose(expected_output, actual_output), "Output mismatch between compiled and original model"