Update Dockerfile.gpu

hopefully wont run out of space
VIDA-NYU · Oct 8, 2024 · 1153d55 · 1153d55
1 parent 36d1729
commit 1153d55
Showing 1 changed file with 2 additions and 3 deletions.
diff --git a/Dockerfile.gpu b/Dockerfile.gpu
@@ -1,4 +1,4 @@
-FROM pytorch/pytorch:1.11.0-cuda11.3-cudnn8-devel
+FROM pytorch/pytorch:2.4.0-cuda11.8-cudnn9-runtime
 RUN rm /etc/apt/sources.list.d/cuda.list
 RUN rm /etc/apt/sources.list.d/nvidia-ml.list
 ARG DEBIAN_FRONTEND=noninteractive
@@ -7,8 +7,7 @@ ENV TZ="America/New_York"
 ENV DIR=/src/lib/ptgctl
 
 RUN apt-get update && apt-get install -y python3-opencv git && rm -rf /var/lib/apt/lists/*
-
-RUN pip install --no-cache-dir -U torch torchvision opencv-python  # for build convenience
+RUN pip install --no-cache-dir -U opencv-python  # for build convenience
 
 ADD setup.py README.md LICENSE $DIR/
 ADD ptgctl/__init__.py ptgctl/__version__.py $DIR/ptgctl/