From c99a16998618d18d8fb27af2f3521677f94861b4 Mon Sep 17 00:00:00 2001 From: jaydrennan Date: Sun, 17 Dec 2023 22:35:12 -0800 Subject: [PATCH] fix: configures test_set_gpu_full to run on a m1 mac. --- tests/test_utils/test_model_cache.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/tests/test_utils/test_model_cache.py b/tests/test_utils/test_model_cache.py index 8b0ed66..1d117e2 100644 --- a/tests/test_utils/test_model_cache.py +++ b/tests/test_utils/test_model_cache.py @@ -65,10 +65,14 @@ def test_set_cpu_full(): @pytest.mark.skipif(get_device() == "cpu", reason="GPU not available") def test_set_gpu_full(): + device = get_device() cache = GPUModelCache( - max_cpu_memory_gb=1, max_gpu_memory_gb=0.0000001, device=get_device() + max_cpu_memory_gb=1, max_gpu_memory_gb=0.0000001, device=device ) - assert cache.max_cpu_memory == 1073741824 + if device in ("cpu", "mps"): + assert cache.max_cpu_memory == 0 + else: + assert cache.max_cpu_memory == 1073741824 model = create_model_of_n_bytes(100_000) with pytest.raises(RuntimeError): cache.set("key1", model)