using image with pytoch pre-installed

alexandriaproject-io · Jan 27, 2024 · e6cc84f · e6cc84f
1 parent 027b71f
commit e6cc84f
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/Dockerfile b/Dockerfile
@@ -1,10 +1,10 @@
-FROM python:3.10-slim
+FROM pytorch/pytorch:2.1.0-cuda11.8-cudnn8-runtime
 
 # Install system dependencies
 RUN apt-get update
 
-# Install PyTorch
-RUN pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118
+# Install PyTorch manually
+# RUN pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118
 
 # Create a cache directory
 RUN mkdir ./cache