fix: libcudart path for some versions of pytorch (#726)

PygmalionAI · Sep 13, 2024 · 9c0e7d9 · 9c0e7d9
1 parent 4648f16
commit 9c0e7d9
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/aphrodite/distributed/device_communicators/cuda_wrapper.py b/aphrodite/distributed/device_communicators/cuda_wrapper.py
@@ -95,8 +95,10 @@ class CudaRTLibrary:
     path_to_dict_mapping: Dict[str, Dict[str, Any]] = {}
 
     def __init__(self, so_file: Optional[str] = None):
-        if so_file is None:
-            so_file = get_pytorch_default_cudart_library_path()
+        if so_file is None: 
+            assert torch.version.cuda is not None 
+            major_version = torch.version.cuda.split(".")[0] 
+            so_file = f"libcudart.so.{major_version}" 
         if so_file not in CudaRTLibrary.path_to_library_cache:
             lib = ctypes.CDLL(so_file)
             CudaRTLibrary.path_to_library_cache[so_file] = lib