fix number devices

huggingface · Mar 12, 2024 · 0089a7c · 0089a7c
1 parent 7ade74f
commit 0089a7c
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/tests/test_api.py b/tests/test_api.py
@@ -44,7 +44,7 @@
 LAUNCHER_CONFIGS = [
     InlineConfig(device_isolation=False),
     ProcessConfig(device_isolation=False),
-    TorchrunConfig(device_isolation=False, nproc_per_node=2),
+    TorchrunConfig(device_isolation=False, nproc_per_node=4),
 ]
 BACKENDS = ["pytorch", "none"]
 DEVICES = ["cpu", "cuda"]
@@ -113,10 +113,12 @@ def test_api_memory_tracker(device, backend):
 @pytest.mark.parametrize("device", DEVICES)
 @pytest.mark.parametrize("launcher_config", LAUNCHER_CONFIGS)
 def test_api_launch(device, launcher_config):
+    device_ids = "0,1,2,3" if device == "cuda" else None
+
     benchmark_config = InferenceConfig(latency=True, memory=True)
     backend_config = PyTorchConfig(
         model="bert-base-uncased",
-        device_ids="0,1" if device == "cuda" else None,
+        device_ids=device_ids,
         no_weights=True,
         device=device,
     )