fix setup.py

setup.py

@@ -19,18 +19,12 @@ SPCONV_FORCE_BUILD_CUDA = os.getenv("SPCONV_FORCE_BUILD_CUDA")
 
 PYTHON_VERSION = "{}.{}".format(sys.version_info.major, sys.version_info.minor)
 
-remove_plus = torch.__version__.find("+dev")
-remove_dot = torch.__version__.find(".dev")
-
+remove_device = re.search(r"(\+|\.)(dev|cu|cpu)", torch.__version__)
 PYTORCH_VERSION = torch.__version__
-if remove_plus != -1:
-    PYTORCH_VERSION = torch.__version__[:remove_plus]
-if remove_dot != -1:
-    PYTORCH_VERSION = torch.__version__[:remove_dot]
-
+if remove_device is not None:
+    PYTORCH_VERSION = torch.__version__[:remove_device.start()]
 PYTORCH_VERSION = list(map(int, PYTORCH_VERSION.split(".")))
 PYTORCH_VERSION_NUMBER = PYTORCH_VERSION[0] * 10000 + PYTORCH_VERSION[1] * 100 + PYTORCH_VERSION[2]
-
 class CMakeExtension(Extension):
     def __init__(self, name, sourcedir='', library_dirs=[]):
         Extension.__init__(self, name, sources=[], library_dirs=library_dirs)