Fix benchmark experiment tests.

pytorch · Oct 1, 2024 · 4f25209 · 4f25209
1 parent 95c3002
commit 4f25209
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/test/benchmarks/test_benchmark_experiment.py b/test/benchmarks/test_benchmark_experiment.py
@@ -7,7 +7,7 @@ class BenchmarkExperimentTest(unittest.TestCase):
 
   def test_to_dict(self):
     be = BenchmarkExperiment("cpu", "PJRT", "some xla_flags", "openxla", None,
-                             False, "train", "123")
+                             False, "train", "123", False)
     actual = be.to_dict()
     self.assertEqual(9, len(actual))
     self.assertEqual("cpu", actual["accelerator"])
@@ -19,6 +19,7 @@ def test_to_dict(self):
     self.assertEqual(False, actual["keep_model_data_on_cuda"])
     self.assertEqual("train", actual["test"])
     self.assertEqual("123", actual["batch_size"])
+    self.assertEqual(False, actual["enable_functionalization"])
 
 
 if __name__ == '__main__':