@@ -18,31 +18,30 @@ export USE_STATIC_NCCL=1
18
18
export ATEN_STATIC_CUDA=1
19
19
export USE_CUDA_STATIC_LINK=1
20
20
21
+ # Keep an array of cmake variables to add to
22
+ if [[ -z " $CMAKE_ARGS " ]]; then
23
+ CMAKE_ARGS=()
24
+ fi
25
+
26
+ # Determine CUDA version and architectures to build for
21
27
CUDA_VERSION=$( nvcc --version| tail -n1| cut -f5 -d" " | cut -f1 -d" ," )
28
+ echo " CUDA $CUDA_VERSION Detected"
22
29
23
30
export TORCH_CUDA_ARCH_LIST=" 3.5;5.0+PTX"
24
31
if [[ $CUDA_VERSION == " 8.0" ]]; then
25
- echo " CUDA 8.0 Detected"
26
32
export TORCH_CUDA_ARCH_LIST=" $TORCH_CUDA_ARCH_LIST ;6.0;6.1"
27
33
elif [[ $CUDA_VERSION == " 9.0" ]]; then
28
- echo " CUDA $CUDA_VERSION Detected"
29
34
export TORCH_CUDA_ARCH_LIST=" $TORCH_CUDA_ARCH_LIST ;6.0;7.0"
30
35
elif [[ $CUDA_VERSION == " 9.2" ]]; then
31
- echo " CUDA $CUDA_VERSION Detected"
32
36
export TORCH_CUDA_ARCH_LIST=" $TORCH_CUDA_ARCH_LIST ;6.0;6.1;7.0"
33
- fi
34
- echo $TORCH_CUDA_ARCH_LIST
35
-
36
- if [[ $CUDA_VERSION == " 8.0" ]]; then
37
- WHEELHOUSE_DIR=" wheelhouse80"
38
- elif [[ $CUDA_VERSION == " 9.0" ]]; then
39
- WHEELHOUSE_DIR=" wheelhouse90"
40
- elif [[ $CUDA_VERSION == " 9.2" ]]; then
41
- WHEELHOUSE_DIR=" wheelhouse92"
37
+ # ATen tests can't build with CUDA 9.2 and the old compiler used here
38
+ CMAKE_ARGS+=(" -DATEN_NO_TEST=ON" )
42
39
else
43
40
echo " unknown cuda version $CUDA_VERSION "
44
41
exit 1
45
42
fi
43
+ echo $TORCH_CUDA_ARCH_LIST
44
+ WHEELHOUSE_DIR=" wheelhouse${CUDA_VERSION: 0: 1} {CUDA_VERSION:2:1}"
46
45
47
46
# Build for given Python versions, or for all in /opt/python if none given
48
47
if [[ -z " $DESIRED_PYTHON " ]]; then
@@ -90,7 +89,7 @@ for PYDIR in "${python_installations[@]}"; do
90
89
else
91
90
pip install numpy==1.11
92
91
fi
93
- time python setup.py bdist_wheel -d $WHEELHOUSE_DIR
92
+ time CMAKE_ARGS= ${CMAKE_ARGS[@]} python setup.py bdist_wheel -d $WHEELHOUSE_DIR
94
93
done
95
94
96
95
popd
0 commit comments