Skip to content

Commit a3ba01a

Browse files
pmeieratalman
andauthored
Use regular env setup (#7926)
Co-authored-by: Andrey Talman <[email protected]>
1 parent fb115c2 commit a3ba01a

File tree

1 file changed

+15
-5
lines changed

1 file changed

+15
-5
lines changed

.github/scripts/setup-env.sh

Lines changed: 15 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -73,11 +73,21 @@ else
7373
CHANNEL=nightly
7474
fi
7575

76-
pip install --progress-bar=off light-the-torch
77-
ltt install --progress-bar=off \
78-
--pytorch-computation-backend="${GPU_ARCH_TYPE}${GPU_ARCH_VERSION}" \
79-
--pytorch-channel="${CHANNEL}" \
80-
torch
76+
case $GPU_ARCH_TYPE in
77+
cpu)
78+
GPU_ARCH_ID="cpu"
79+
;;
80+
cuda)
81+
VERSION_WITHOUT_DOT=$(echo "${GPU_ARCH_VERSION}" | sed 's/\.//')
82+
GPU_ARCH_ID="cu${VERSION_WITHOUT_DOT}"
83+
;;
84+
*)
85+
echo "Unknown GPU_ARCH_TYPE=${GPU_ARCH_TYPE}"
86+
exit 1
87+
;;
88+
esac
89+
PYTORCH_WHEEL_INDEX="https://download.pytorch.org/whl/${CHANNEL}/${GPU_ARCH_ID}"
90+
pip install --progress-bar=off --pre torch --index-url="${PYTORCH_WHEEL_INDEX}"
8191

8292
if [[ $GPU_ARCH_TYPE == 'cuda' ]]; then
8393
python -c "import torch; exit(not torch.cuda.is_available())"

0 commit comments

Comments
 (0)