From 2296573714b35b223d38735c1ee9a0e4b5dfbfd6 Mon Sep 17 00:00:00 2001 From: Jinzhe Zeng Date: Thu, 4 Jan 2024 22:55:16 -0500 Subject: [PATCH] fix CPU tests on GPUs Signed-off-by: Jinzhe Zeng --- source/tests/test_auto_batch_size.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/source/tests/test_auto_batch_size.py b/source/tests/test_auto_batch_size.py index 93a96c9c29..5a349f70b9 100644 --- a/source/tests/test_auto_batch_size.py +++ b/source/tests/test_auto_batch_size.py @@ -45,8 +45,10 @@ def test_execute_oom_gpu(self, mock_is_gpu_available): self.assertEqual(result.shape, (256, 2)) @unittest.mock.patch("tensorflow.compat.v1.test.is_gpu_available") - def test_execute_oom_cpu(self, mock_is_gpu_available): + @unittest.mock.patch("tensorflow.compat.v1.config.experimental.get_visible_devices") + def test_execute_oom_cpu(self, mock_is_gpu_available, mock_get_visible_devices): mock_is_gpu_available.return_value = False + mock_get_visible_devices.return_value = [] # initial batch size 256 = 128 * 2, nb is always 128 auto_batch_size = AutoBatchSize(256, 2.0) nb, result = auto_batch_size.execute(self.oom, 1, 2)