fixup

vllm-project · May 10, 2024 · bdf6378 · bdf6378
1 parent 62cde18
commit bdf6378
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/vllm/_custom_ops.py b/vllm/_custom_ops.py
@@ -5,11 +5,12 @@
 try:
     from vllm._C import cache_ops as vllm_cache_ops
     from vllm._C import ops as vllm_ops
+
+    capability = torch.cuda.get_device_capability()
+    capability = capability[0] * 10 + capability[1]
 except ImportError:
     pass
 
-capability = torch.cuda.get_device_capability()
-capability = capability[0] * 10 + capability[1]
 
 
 # activation ops