From eb931c7f88fc055da911f9fbc8535abf6bd68925 Mon Sep 17 00:00:00 2001 From: Cathy Wong Date: Wed, 2 Mar 2022 13:25:35 -0500 Subject: [PATCH] [MD] test_pyfunc_multiprocess_autotune.py - enable AutoTune --- .../ut/python/dataset/test_pyfunc_multiprocess_autotune.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/tests/ut/python/dataset/test_pyfunc_multiprocess_autotune.py b/tests/ut/python/dataset/test_pyfunc_multiprocess_autotune.py index 59eebb705f0..0f7f075a91a 100644 --- a/tests/ut/python/dataset/test_pyfunc_multiprocess_autotune.py +++ b/tests/ut/python/dataset/test_pyfunc_multiprocess_autotune.py @@ -138,6 +138,10 @@ class TestPythonMultiprocAutotune: """ Run before each test function. """ + # Enable Dataset AutoTune + self.original_autotune = ds.config.get_enable_autotune() + ds.config.set_enable_autotune(True) + # Reduce memory required by disabling the shared memory optimization self.mem_original = ds.config.get_enable_shared_mem() ds.config.set_enable_shared_mem(False) @@ -146,7 +150,9 @@ class TestPythonMultiprocAutotune: """ Run after each test function. """ + # Restore settings ds.config.set_enable_shared_mem(self.mem_original) + ds.config.set_enable_autotune(self.original_autotune) @staticmethod def test_cifar10_pyfunc_pipeline():