fix

Lightning-AI · Jun 15, 2024 · 8ff1dcd · 8ff1dcd
1 parent a0086f1
commit 8ff1dcd
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/src/litserve/connector.py b/src/litserve/connector.py
@@ -72,6 +72,9 @@ def _auto_device_count(self, accelerator) -> int:
 
     @staticmethod
     def _choose_gpu_accelerator_backend():
+        if check_cuda_with_nvidia_smi() > 0:
+            return "cuda"
+
         try:
             import torch
 
@@ -80,9 +83,6 @@ def _choose_gpu_accelerator_backend():
         except ImportError:
             return None
 
-        if check_cuda_with_nvidia_smi() > 0:
-            return "cuda"
-
         return None