| | #!/bin/bash -e |
| | |
| |
|
| | |
| | retry () { |
| | $* || (sleep 1 && $*) || (sleep 2 && $*) || (sleep 4 && $*) || (sleep 8 && $*) |
| | } |
| | |
| | pip_install() { |
| | retry pip install --progress-bar off "$@" |
| | } |
| |
|
| |
|
| | setup_cuda() { |
| | |
| | |
| | export FORCE_CUDA=1 |
| | case "$CU_VERSION" in |
| | cu102) |
| | export CUDA_HOME=/usr/local/cuda-10.2/ |
| | export TORCH_CUDA_ARCH_LIST="3.5;3.7;5.0;5.2;6.0+PTX;6.1+PTX;7.0+PTX;7.5+PTX" |
| | ;; |
| | cu101) |
| | export CUDA_HOME=/usr/local/cuda-10.1/ |
| | export TORCH_CUDA_ARCH_LIST="3.5;3.7;5.0;5.2;6.0+PTX;6.1+PTX;7.0+PTX;7.5+PTX" |
| | ;; |
| | cu100) |
| | export CUDA_HOME=/usr/local/cuda-10.0/ |
| | export TORCH_CUDA_ARCH_LIST="3.5;3.7;5.0;5.2;6.0+PTX;6.1+PTX;7.0+PTX;7.5+PTX" |
| | ;; |
| | cu92) |
| | export CUDA_HOME=/usr/local/cuda-9.2/ |
| | export TORCH_CUDA_ARCH_LIST="3.5;3.7;5.0;5.2;6.0+PTX;6.1+PTX;7.0+PTX" |
| | ;; |
| | cpu) |
| | unset FORCE_CUDA |
| | export CUDA_VISIBLE_DEVICES= |
| | ;; |
| | *) |
| | echo "Unrecognized CU_VERSION=$CU_VERSION" |
| | exit 1 |
| | ;; |
| | esac |
| | } |
| |
|
| | setup_wheel_python() { |
| | case "$PYTHON_VERSION" in |
| | 3.6) python_abi=cp36-cp36m ;; |
| | 3.7) python_abi=cp37-cp37m ;; |
| | 3.8) python_abi=cp38-cp38 ;; |
| | *) |
| | echo "Unrecognized PYTHON_VERSION=$PYTHON_VERSION" |
| | exit 1 |
| | ;; |
| | esac |
| | export PATH="/opt/python/$python_abi/bin:$PATH" |
| | } |
| |
|