make sure testing model doesn't go onto MPS (insufficient memory on g…

…h's shared setup)
NVIDIA · Aug 1, 2024 · 5a2f100 · 5a2f100
1 parent 4d616fb
commit 5a2f100
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/tests/probes/test_probes_fileformats.py b/tests/probes/test_probes_fileformats.py
@@ -18,7 +18,9 @@ def test_hf_files_load():
 
 def test_hf_files_hf_repo():
     p = garak._plugins.load_plugin("probes.fileformats.HF_Files")
-    garak._config.plugins.generators["huggingface"] = {"Model": {"name": "gpt2"}}
+    garak._config.plugins.generators["huggingface"] = {
+        "Model": {"name": "gpt2", "hf_args": {"device": "cpu"}},
+    }
     g = garak._plugins.load_plugin(
         "generators.huggingface.Model", config_root=garak._config
     )