Update target_detector.py

microsoft · Jun 30, 2024 · 2e52552 · 2e52552 · LeiWang1999 · Jun 30, 2024
1 parent e1fa655
commit 2e52552
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/python/bitblas/utils/target_detector.py b/python/bitblas/utils/target_detector.py
@@ -44,7 +44,7 @@ def get_gpu_model_from_nvidia_smi(gpu_id: int = 0):
 
     # for multiple gpus, CUDA_DEVICE_ORDER=PCI_BUS_ID must be set to match nvidia-smi or else wrong
     # gpu is returned for gpu_id
-    if len(gpus) > 0 and os.environ.get("CUDA_DEVICE_ORDER") != "PCI_BUS_ID":
+    if len(gpus) > 1 and os.environ.get("CUDA_DEVICE_ORDER") != "PCI_BUS_ID":
         raise EnvironmentError("Multi-gpu environment must set `CUDA_DEVICE_ORDER=PCI_BUS_ID`.")
 
     if gpu_id >= len(gpus) or gpu_id < 0: