I am trying to deploy a pretrained PyTorch model on Google Cloud Platform (GCP). The model works fine on a CPU, but when I add GPU support to the service, the GPU memory is not utilized. Upon investigation, I found that I am using a python:3.10 base image, and I realize that I should be using a CUDA-compatible base image for GPU support.
Currently, I am installing PyTorch and torchvision using pip, but the requirements.txt I received from the pretrained model has the following comment:
torch==1.13.1+cu116 # Kept as comment for version reference
However, when building the service, I encounter an error saying that this specific version is not available.
Can you suggest a lightweight GPU-compatible base image that includes PyTorch and torchvision versions with GPU support? I am looking for a base image that will work with CUDA on GCP.