fix is_fp8_supported function

predibase · Oct 28, 2024 · 813b64b · 813b64b
1 parent 488d428
commit 813b64b
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/server/lorax_server/utils/torch_utils.py b/server/lorax_server/utils/torch_utils.py
@@ -15,8 +15,9 @@ def is_quantized(quantize):
 
 
 def is_fp8_supported():
-    return (torch.cuda.get_device_capability()[0] >= 9) \
-        or (torch.cuda.get_device_capability()[0] == 8 and torch.cuda.get_device_capability()[1] >= 9)
+    return torch.cuda.is_available() and \
+        (torch.cuda.get_device_capability()[0] >= 9) or \
+        (torch.cuda.get_device_capability()[0] == 8 and torch.cuda.get_device_capability()[1] >= 9)
 
 
 def is_fp8_kv(quantize):