Update Dockerfile

Dockerfile

@@ -35,10 +35,6 @@ RUN apt-get update && apt-get install -y --no-install-recommends \
 RUN pip install --upgrade pip setuptools wheel packaging ninja
 
 # Uninstall any prebuilt flash-attn
-RUN pip uninstall -y flash-attn || true
-
-# Build from source so it links correctly
-RUN pip install flash-attn --no-build-isolation --no-cache-dir
 
 ENV CONDA_FORGE=1 \
     CONDA_AUTO_UPDATE_CONDA=false
@@ -56,6 +52,10 @@ RUN pip install --upgrade pip && \
 
 # Install the rest of the dependencies
 RUN pip install --no-cache-dir -r requirements.txt
+RUN pip uninstall -y flash-attn || true
+
+# Build from source so it links correctly
+RUN pip install flash-attn --no-build-isolation --no-cache-dir
 
 # Copy project code
 COPY . .