Install CUDA PyTorch in slim image for A100 GPU

Dockerfile

@@ -10,12 +10,16 @@ ENV XDG_CACHE_HOME=/tmp/cache
 ENV HOME=/tmp/home
 ENV USER=user
 ENV PYTHONUNBUFFERED=1
+ENV NVIDIA_VISIBLE_DEVICES=all
+ENV NVIDIA_DRIVER_CAPABILITIES=compute,utility
 
 RUN mkdir -p /tmp/hf_cache /tmp/torch_cache /tmp/cache /tmp/home && \
     chmod -R 777 /tmp/hf_cache /tmp/torch_cache /tmp/cache /tmp/home
 
 RUN pip install --no-cache-dir \
-    torch \
+    torch --index-url https://download.pytorch.org/whl/cu121
+
+RUN pip install --no-cache-dir \
     transformers \
     trl \
     peft \