From 4f25209a48bf2b7aa9139c82c7705672267424f8 Mon Sep 17 00:00:00 2001 From: Yukio Siraichi Date: Tue, 1 Oct 2024 11:06:32 -0300 Subject: [PATCH] Fix benchmark experiment tests. --- test/benchmarks/test_benchmark_experiment.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/test/benchmarks/test_benchmark_experiment.py b/test/benchmarks/test_benchmark_experiment.py index afc3419e8b8..1b81573ad6e 100644 --- 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__':