v3.0.1

BBC-Esq · Jul 1, 2024 · 68d24d9 · 68d24d9
1 parent 635548e
commit 68d24d9
Showing 1 changed file with 9 additions and 8 deletions.
diff --git a/ct2_main.py b/ct2_main.py
@@ -7,21 +7,22 @@
 from ct2_utils import CheckQuantizationSupport
 
 def set_cuda_paths():
-    nvidia_base_path = Path(sys.executable).parent / 'Lib' / 'site-packages' / 'nvidia'
-    cublas_bin_path = Path(sys.executable).parent / 'Lib' / 'site-packages' / 'nvidia' / 'cublas' / 'bin'
-    cudnn_bin_path = Path(sys.executable).parent / 'Lib' / 'site-packages' / 'nvidia' / 'cudnn' / 'bin'
+    script_dir = Path(__file__).parent.resolve()
+    nvidia_base_path = script_dir / 'Lib' / 'site-packages' / 'nvidia'
+    cublas_bin_path = script_dir / 'Lib' / 'site-packages' / 'nvidia' / 'cublas' / 'bin'
+    cudnn_bin_path = script_dir / 'Lib' / 'site-packages' / 'nvidia' / 'cudnn' / 'bin'
 
     # Set CUDA_PATH and CUDA_PATH_V12_2
     for env_var in ['CUDA_PATH', 'CUDA_PATH_V12_2']:
         current_path = os.environ.get(env_var, '')
         new_paths = [str(nvidia_base_path), str(cublas_bin_path), str(cudnn_bin_path)]
         os.environ[env_var] = os.pathsep.join(filter(None, new_paths + [current_path]))
 
-    # # Add nvidia folder, cudnn bin folder, and cublas bin folder to system PATH
-    # current_path = os.environ.get('PATH', '')
-    # new_paths = [str(nvidia_base_path), str(cublas_bin_path), str(cudnn_bin_path)]
-    # new_path = os.pathsep.join(filter(None, new_paths + [current_path]))
-    # os.environ['PATH'] = new_path
+    # Add nvidia folder, cudnn bin folder, and cublas bin folder to system PATH
+    current_path = os.environ.get('PATH', '')
+    new_paths = [str(nvidia_base_path), str(cublas_bin_path), str(cudnn_bin_path)]
+    new_path = os.pathsep.join(filter(None, new_paths + [current_path]))
+    os.environ['PATH'] = new_path
 
 if __name__ == "__main__":
     set_cuda_paths()