Check the versions of PyTorch and CUDA at the import time. (#949)

* Check the versions of PyTorch and CUDA at the import time.

.flake8

@@ -14,6 +14,7 @@ exclude =
   get_version.py
   build,
   k2/python/host,
+  k2/python/k2/__init__.py,
   k2/python/k2/ctc_loss.py,
   docs
 

.gitignore

@@ -1,3 +1,4 @@
+k2/python/k2/torch_version.py
 # Build folder
 **/build*
 

CMakeLists.txt

@@ -264,6 +264,10 @@ if(K2_USE_PYTORCH)
   add_definitions(-DK2_USE_PYTORCH)
   add_definitions(-DTORCH_API_INCLUDE_EXTENSION_H)
   include(torch)
+  configure_file(
+    ${CMAKE_SOURCE_DIR}/k2/python/k2/torch_version.py.in
+    ${CMAKE_SOURCE_DIR}/k2/python/k2/torch_version.py @ONLY
+  )
 endif()
 
 if(K2_WITH_CUDA)

k2/python/k2/__init__.py

@@ -1,9 +1,29 @@
 import torch  # noqa
+from .torch_version import k2_torch_cuda_version
+from .torch_version import k2_torch_version
+
+if torch.__version__.split("+")[0] != k2_torch_version.split("+")[0]:
+    raise ImportError(
+        f"k2 was built using PyTorch {k2_torch_version}\n"
+        f"But you are using PyTorch {torch.__version__} to run it"
+    )
+
+if (
+    k2_torch_cuda_version != ""
+    and torch.version.cuda is not None
+    and torch.version.cuda != k2_torch_cuda_version
+):
+    raise ImportError(
+        f"k2 was built using CUDA {k2_torch_cuda_version}\n"
+        f"But you are using CUDA {torch.version.cuda} to run it."
+    )
+
 try:
     from _k2 import DeterminizeWeightPushingType
     from _k2 import simple_ragged_index_select
 except ImportError as e:
     import sys
+
     major_v, minor_v = sys.version_info[:2]
     raise ImportError(
         str(e) + "\nNote: If you're using anaconda and importing k2 on MacOS,"
@@ -18,6 +38,7 @@
 from . import dense_fsa_vec
 from . import fsa
 from . import utils
+
 #
 from .autograd import intersect_dense
 from .autograd import intersect_dense_pruned

k2/python/k2/torch_version.py.in

@@ -0,0 +1,17 @@
+# Auto generated by the toplevel CMakeLists.txt.
+#
+# DO NOT EDIT.
+
+# The torch version used to build k2. We will check it against the torch version
+# that is used to run k2. If they are not the same, `import k2` will throw.
+#
+# Some example values are:
+# - 1.10.0+cu102
+# - 1.5.0+cpu
+k2_torch_version = "@TORCH_VERSION@"
+
+# The CUDA version used to build k2.
+# Note: It is an empty string if you used a CPU version of PyTorch to build k2
+#
+# An example value is "10.2".
+k2_torch_cuda_version = "@TORCH_CUDA_VERSION@"