fix batch size

undertherain · Jun 17, 2020 · fca757b · fca757b
1 parent 12b414f
commit fca757b
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/test/pytorch/test_bert.py b/test/pytorch/test_bert.py
@@ -15,8 +15,8 @@ def test_bert(self):
             self.name,
             "--framework=pytorch",
             "--problem=bert",
-            "--problem_size=2,8",
-            "--batch_size=8",
+            "--problem_size=4,8",
+            "--batch_size=2",
             "--nb_epoch=1",
             "--mode=inference",
         )