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 73
73
CHANNEL=nightly
74
74
fi
75
75
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} "
81
91
82
92
if [[ $GPU_ARCH_TYPE == ' cuda' ]]; then
83
93
python -c " import torch; exit(not torch.cuda.is_available())"
You can’t perform that action at this time.
0 commit comments