Spaces:
Sleeping
Sleeping
# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved | |
set -ex | |
ldconfig # https://github.com/NVIDIA/nvidia-docker/issues/854 | |
script_dir="$( cd "$( dirname "${BASH_SOURCE[0]}" )" >/dev/null 2>&1 && pwd )" | |
. "$script_dir/pkg_helpers.bash" | |
echo "Build Settings:" | |
echo "CU_VERSION: $CU_VERSION" # e.g. cu101 | |
echo "D2_VERSION_SUFFIX: $D2_VERSION_SUFFIX" # e.g. +cu101 or "" | |
echo "PYTHON_VERSION: $PYTHON_VERSION" # e.g. 3.6 | |
echo "PYTORCH_VERSION: $PYTORCH_VERSION" # e.g. 1.4 | |
setup_cuda | |
setup_wheel_python | |
yum install ninja-build -y && ln -sv /usr/bin/ninja-build /usr/bin/ninja | |
export TORCH_VERSION_SUFFIX="+$CU_VERSION" | |
if [[ "$CU_VERSION" == "cu102" ]]; then | |
export TORCH_VERSION_SUFFIX="" | |
fi | |
pip_install pip numpy -U | |
pip_install "torch==$PYTORCH_VERSION$TORCH_VERSION_SUFFIX" \ | |
-f https://download.pytorch.org/whl/$CU_VERSION/torch_stable.html | |
# use separate directories to allow parallel build | |
BASE_BUILD_DIR=build/$CU_VERSION/$PYTHON_VERSION | |
python setup.py \ | |
build -b $BASE_BUILD_DIR \ | |
bdist_wheel -b $BASE_BUILD_DIR/build_dist -d wheels/$CU_VERSION | |