Set num_devices in Python example

Cambridge-ICCS · Feb 6, 2025 · 3b67ae8 · 3b67ae8
1 parent 3793c3d
commit 3b67ae8
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/examples/3_MultiGPU/multigpu_infer_python.py b/examples/3_MultiGPU/multigpu_infer_python.py
@@ -53,7 +53,9 @@ def deploy(saved_model: str, device: str, batch_size: int = 1) -> torch.Tensor:
 if __name__ == "__main__":
     saved_model_file = "saved_multigpu_model_cuda.pt"
 
-    for device_index in range(2):
+    num_devices = 2
+
+    for device_index in range(num_devices):
         device_to_run = f"cuda:{device_index}"
 
         batch_size_to_run = 1