From 15809ab36ccbd29920ca4d084c0c6e8071212533 Mon Sep 17 00:00:00 2001 From: Soumith Chintala Date: Fri, 22 Feb 2019 14:20:52 -0800 Subject: [PATCH] smoke test fix --- smoke_test.sh | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/smoke_test.sh b/smoke_test.sh index 6bacd16e5..e909f22f9 100755 --- a/smoke_test.sh +++ b/smoke_test.sh @@ -81,7 +81,8 @@ if [[ "$PACKAGE_TYPE" == 'conda' ]]; then rm -rf /usr/local/cuda || true if [[ "$DESIRED_CUDA" != 'cpu' ]]; then ln -s "/usr/local/cuda-${cuda_dot}" /usr/local/cuda - export CUDA_VERSION=$(ls /usr/local/cuda/lib64/libcudart.so.*|sort|tac | head -1 | rev | cut -d"." -f -3 | rev) + export CUDA_VERSION=$(ls /usr/local/cuda/lib64/libcudart.so.*|sort|tac | head -1 | rev | cut -d"." -f -3 | rev) # 10.0.130 + export CUDA_VERSION_SHORT=$(ls /usr/local/cuda/lib64/libcudart.so.*|sort|tac | head -1 | rev | cut -d"." -f -3 | rev | cut -f1,2 -d".") # 10.0 export CUDNN_VERSION=$(ls /usr/local/cuda/lib64/libcudnn.so.*|sort|tac | head -1 | rev | cut -d"." -f -3 | rev) fi fi @@ -103,10 +104,10 @@ if [[ "$PACKAGE_TYPE" == 'libtorch' ]]; then retry curl -o libtorch_zip "https://download.pytorch.org/libtorch/nightly/$DESIRED_CUDA/$package_name" unzip -q libtorch_zip elif [[ "$PACKAGE_TYPE" == 'conda' ]]; then - if [[ "$DESIRED_CUDA" == 'cpu' || "$DESIRED_CUDA" == 'cu90' ]]; then + if [[ "$DESIRED_CUDA" == 'cpu' ]]; then retry conda install -yq -c pytorch "$package_name_and_version" else - retry conda install -yq -c pytorch "cuda${DESIRED_CUDA:2}" "$package_name_and_version" + retry conda install -yq -c pytorch "cudatoolkit=$CUDA_VERSION_SHORT" "$package_name_and_version" fi else retry pip install "$package_name_and_version" \