| # 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. | |