File size: 258 Bytes
db12eaa
 
 
 
1
2
3
4
5
# torch, CUDA 12.1, and cuDNN 8 are pre-installed in the base image:
# pytorch/pytorch:2.3.0-cuda12.1-cudnn8-runtime
# Do NOT add torch here — pip would resolve to the CPU wheel from default PyPI
# and overwrite the CUDA-enabled torch from the base image.