update

Author: rusty1s

script/cuda.sh

@@ -95,7 +95,7 @@ if [ "${TRAVIS_OS_NAME}" = "linux" ] && [ "${IDX}" != "cpu" ] && [ "${IDX}" != "
   nvcc --version
 fi
 
-if [ "${TRAVIS_OS_NAME}" = "linux" ] [ "${IDX}" = "cu110" ]; then
+if [ "${TRAVIS_OS_NAME}" = "linux" ] && [ "${IDX}" = "cu110" ]; then
   wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu1804/x86_64/cuda-ubuntu1804.pin
   sudo mv cuda-ubuntu1804.pin /etc/apt/preferences.d/cuda-repository-pin-600
   wget https://developer.download.nvidia.com/compute/cuda/11.0.3/local_installers/cuda-repo-ubuntu1804-11-0-local_11.0.3-450.51.06-1_amd64.deb

script/torch.sh

@@ -7,16 +7,16 @@ if [ "${TRAVIS_OS_NAME}" = "windows" ]; then
   sed -i.bak -e 's/return \*(this->value)/return \*((type\*)this->value)/g' /c/tools/miniconda3/envs/test/lib/site-packages/torch/include/pybind11/cast.h
 fi
 
-if [ "${TRAVIS_OS_NAME}" = "windows" ] && [ "${TORCH_VERSION}" = "1.7.0" ]; then
-  echo "Fix nvcc for PyTorch 1.7.0"
-  sed -i.bak -e 's/constexpr Symbol Kind = ::c10::prim::profile;/const Symbol Kind;/g' /c/tools/miniconda3/envs/test/lib/site-packages/torch/include/torch/csrc/jit/ir/ir.h
-  sed -i.bak '1345a\
-const Symbol ProfileOp::Kind = ::c10::prim::profile;' /c/tools/miniconda3/envs/test/lib/site-packages/torch/include/torch/csrc/jit/ir/ir.h
-  sed -i.bak -e 's/constexpr Symbol Kind = ::c10::prim::profile_optional;/const Symbol Kind;/g' /c/tools/miniconda3/envs/test/lib/site-packages/torch/include/torch/csrc/jit/ir/ir.h
-  sed -i.bak '1368a\
-const Symbol ProfileOptionalOp::Kind = ::c10::prim::profile_optional;' /c/tools/miniconda3/envs/test/lib/site-packages/torch/include/torch/csrc/jit/ir/ir.h
+# if [ "${TRAVIS_OS_NAME}" = "windows" ] && [ "${TORCH_VERSION}" = "1.7.0" ]; then
+#   echo "Fix nvcc for PyTorch 1.7.0"
+#   sed -i.bak -e 's/constexpr Symbol Kind = ::c10::prim::profile;/const Symbol Kind;/g' /c/tools/miniconda3/envs/test/lib/site-packages/torch/include/torch/csrc/jit/ir/ir.h
+#   sed -i.bak '1345a\
+# const Symbol ProfileOp::Kind = ::c10::prim::profile;' /c/tools/miniconda3/envs/test/lib/site-packages/torch/include/torch/csrc/jit/ir/ir.h
+#   sed -i.bak -e 's/constexpr Symbol Kind = ::c10::prim::profile_optional;/const Symbol Kind;/g' /c/tools/miniconda3/envs/test/lib/site-packages/torch/include/torch/csrc/jit/ir/ir.h
+#   sed -i.bak '1368a\
+# const Symbol ProfileOptionalOp::Kind = ::c10::prim::profile_optional;' /c/tools/miniconda3/envs/test/lib/site-packages/torch/include/torch/csrc/jit/ir/ir.h
 
-  echo "-------------------------"
-  sed -n '1327,1375p' /c/tools/miniconda3/envs/test/lib/site-packages/torch/include/torch/csrc/jit/ir/ir.h
-  echo "-------------------------"
-fi
+#   echo "-------------------------"
+#   sed -n '1327,1375p' /c/tools/miniconda3/envs/test/lib/site-packages/torch/include/torch/csrc/jit/ir/ir.h
+#   echo "-------------------------"
+# fi

setup.py

@@ -1,5 +1,6 @@
 import os
 import os.path as osp
+import sys
 import glob
 from setuptools import setup, find_packages
 
@@ -26,7 +27,11 @@ def get_extensions():
     info = parallel_info()
     print(info)
     if 'parallel backend: OpenMP' in info and 'OpenMP not found' not in info:
-        extra_compile_args['cxx'] += ['-DAT_PARALLEL_OPENMP', '-fopenmp']
+        extra_compile_args['cxx'] += ['-DAT_PARALLEL_OPENMP']
+        if sys.platform == 'win32':
+            extra_compile_args['cxx'] += ['/openmp']
+        else:
+            extra_compile_args['cxx'] += ['-fopenmp']
     else:
         print('Compiling without OpenMP...')