File tree Expand file tree Collapse file tree 1 file changed +15
-5
lines changed Expand file tree Collapse file tree 1 file changed +15
-5
lines changed Original file line number Diff line number Diff line change 7373 CHANNEL=nightly
7474fi
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
8292if [[ $GPU_ARCH_TYPE == ' cuda' ]]; then
8393 python -c " import torch; exit(not torch.cuda.is_available())"
You can’t perform that action at this time.
0 commit comments