Merge 9fba12c into 21b87dc

undertherain · Jun 19, 2020 · e14ac9d · e14ac9d
2 parents 21b87dc + 9fba12c
commit e14ac9d
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 2 deletions.
diff --git a/benchmarker/modules/problems/bert/data.py b/benchmarker/modules/problems/bert/data.py
@@ -10,5 +10,5 @@ def get_data(params):
              params["problem"]["len_sequence"],
              )
     X = np.random.random(shape).astype(np.int64)
-    Y = np.ones((cnt_batches, params["batch_size"]))
+    Y = np.ones((cnt_batches, params["batch_size"]), dtype=np.int64)
     return X, Y
diff --git a/test/pytorch/test_bert.py b/test/pytorch/test_bert.py
@@ -10,7 +10,7 @@ class PytorchBertTest(unittest.TestCase):
     def setUp(self):
         self.name = "benchmarker"
 
-    def test_bert(self):
+    def test_bert_inference(self):
         run_module(
             self.name,
             "--framework=pytorch",
@@ -20,3 +20,14 @@ def test_bert(self):
             "--nb_epoch=1",
             "--mode=inference",
         )
+
+    def test_bert_training(self):
+        run_module(
+            self.name,
+            "--framework=pytorch",
+            "--problem=bert",
+            "--problem_size=4,8",
+            "--batch_size=2",
+            "--nb_epoch=1",
+            "--mode=training",
+        )