Fix building pip packages for CUDA 11.0 (#584)

* fix building pip packages for CUDA 11.0
k2-fsa · Jan 12, 2021 · fe00848 · fe00848
1 parent c2a6f9c
commit fe00848
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 3 deletions.
diff --git a/CMakeLists.txt b/CMakeLists.txt
@@ -13,7 +13,7 @@ cmake_minimum_required(VERSION 3.8 FATAL_ERROR)
 
 project(k2 CUDA CXX)
 
-set(K2_VERSION "0.1.2")
+set(K2_VERSION "0.1.3")
 
 # ----------------- Supported build types for K2 project -----------------
 set(ALLOWABLE_BUILD_TYPES Debug Release RelWithDebInfo MinSizeRel)

diff --git a/setup.py b/setup.py
@@ -39,9 +39,15 @@ def get_long_description():
 
 
 def get_cuda_version():
+    import torch
     from torch.utils import collect_env
-    cuda_version = collect_env.get_running_cuda_version(
-        collect_env.run).split('.')
+    running_cuda_version = collect_env.get_running_cuda_version(collect_env.run)
+    cuda_version = torch.version.cuda
+    if running_cuda_version is not None:
+        assert cuda_version in running_cuda_version, \
+                f'PyTorch is built with CUDA version: {cuda_version}.\n' \
+                f'The current running CUDA version is: {running_cuda_version}'
+    cuda_version = cuda_version.split('.')
     major, minor = int(cuda_version[0]), int(cuda_version[1])
     cuda_version = major * 10 + minor
     return f'{cuda_version}'