torch version string man fix

NVIDIA · Nov 15, 2021 · a7099e2 · a7099e2
1 parent 12663c7
commit a7099e2
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/setup.py b/setup.py
@@ -147,7 +147,7 @@ def _argparse(pattern, argv, is_flag=True, is_list=False):
 CC_FLAGS = []
 NVCC_FLAGS = []
 
-TORCH_VERSION = [int(v) for v in torch.__version__.split('.')]
+TORCH_VERSION = [int(v) for v in torch.__version__.replace('+', '.').split('.')[:3]]
 
 if CPU_ONLY:
     print("--------------------------------")