Update Dockerfile

Dockerfile

@@ -3,7 +3,7 @@ FROM nvidia/cuda:11.8.0-cudnn8-devel-ubuntu20.04
 
 ENV DEBIAN_FRONTEND=noninteractive
 
-# Basic utilities
+# Install basic utilities
 RUN apt-get update && apt-get install -y \
     tzdata \
     software-properties-common \
@@ -19,10 +19,13 @@ RUN add-apt-repository ppa:deadsnakes/ppa && apt-get update && \
 RUN curl -sS https://bootstrap.pypa.io/get-pip.py | python3.11
 RUN update-alternatives --install /usr/bin/python python /usr/bin/python3.11 1
 
-# Install Python packages (torch, Transformers, Accelerate, etc.)
+# Install Python packages (torch, xFormers, Transformers, etc.)
 RUN python -m pip install --no-cache-dir \
     torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118 && \
     python -m pip install --no-cache-dir \
+    xformers==0.0.27 \
+        --extra-index-url https://download.pytorch.org/whl/nightly/cu118/torch_nightly.html && \
+    python -m pip install --no-cache-dir \
     transformers \
     accelerate \
     trl \