cuda 10.1 fixups

Signed-off-by: Edward Z. Yang <[email protected]>

Author: ezyang

packaging/pkg_helpers.bash

@@ -52,6 +52,13 @@ setup_cuda() {
 
   # Now work out the CUDA settings
   case "$CU_VERSION" in
+    cu101)
+      export CUDA_HOME=/usr/local/cuda-10.1/
+      export FORCE_CUDA=1
+      # Hard-coding gencode flags is temporary situation until
+      # https://github.com/pytorch/pytorch/pull/23408 lands
+      export NVCC_FLAGS="-gencode=arch=compute_35,code=sm_35 -gencode=arch=compute_50,code=sm_50 -gencode=arch=compute_60,code=sm_60 -gencode=arch=compute_70,code=sm_70 -gencode=arch=compute_75,code=sm_75 -gencode=arch=compute_50,code=compute_50"
+      ;;
     cu100)
       export CUDA_HOME=/usr/local/cuda-10.0/
       export FORCE_CUDA=1

packaging/wheel/linux_manywheel.sh

@@ -8,7 +8,7 @@ if [ "$#" -ne 1 ]; then
 fi
 export CUVER="$1" # cu92 cu100 cpu
 
-if [[ "$CUVER" == "cu100" ]]; then
+if [[ "$CUVER" == "cu101" ]]; then
   cu_suffix=""
 else
   cu_suffix="+$CUVER"