Update Dockerfile

Dockerfile

@@ -46,11 +46,14 @@ WORKDIR /workspace
 COPY requirements.txt .
 
 # Pre-install torch so flash-attn doesn't fail
-RUN pip install --upgrade pip && \
-    pip install torch==2.6 torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu121
-    
+# RUN pip install --upgrade pip && \
+#    pip install torch==2.6 torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu121
+
+RUN pip install torch==2.5.1+cu121 torchvision --index-url https://download.pytorch.org/whl/cu121
+RUN pip install flash-attn
+
 # Build flash-attn from source (matches PyTorch ABI)
-RUN pip install --no-build-isolation --force-reinstall flash-attn
+# RUN pip install --no-build-isolation --force-reinstall flash-attn
 
 RUN pip install -r requirements.txt
 # Copy project code