torch 116

Dockerfile

-FROM nvidia/cuda:11.3.1-base-ubuntu20.04
+FROM nvidia/cuda:11.6.1-base-ubuntu20.04
 ENV DEBIAN_FRONTEND=noninteractive
 
 #Install core packages
@@ -11,7 +11,7 @@ RUN apt-get update && apt-get install -y \
     pip3 install --no-cache-dir --upgrade pip
 
 #Install Torch(chonky!)
-RUN pip3 install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu113
+RUN pip3 install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu116
 
 #Install Python deps
 RUN pip3 install --no-cache-dir dotmap icecream sentry-sdk numpy fastapi "uvicorn[standard]" gunicorn