make similar changes as in rapidsai#4436 to resolve wheel test failures

trxcllnt · May 24, 2024 · c8f66e4 · c8f66e4
1 parent f7a54c9
commit c8f66e4
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/ci/test_wheel_cugraph-dgl.sh b/ci/test_wheel_cugraph-dgl.sh
@@ -33,7 +33,10 @@ PYTORCH_URL="https://download.pytorch.org/whl/cu${PYTORCH_CUDA_VER}"
 DGL_URL="https://data.dgl.ai/wheels/cu${PYTORCH_CUDA_VER}/repo.html"
 
 rapids-logger "Installing PyTorch and DGL"
-rapids-retry python -m pip install torch --index-url ${PYTORCH_URL}
+rapids-retry python -m pip install --no-cache-dir torch --index-url ${PYTORCH_URL}
 rapids-retry python -m pip install dgl==2.0.0 --find-links ${DGL_URL}
 
+python -m pip uninstall -y torch torchvision torchaudio
+python -m pip install --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/cu118
+
 python -m pytest python/cugraph-dgl/tests