From fa96d6467ad93c37d8d5569f3cade435bf7c46c6 Mon Sep 17 00:00:00 2001
From: abhishek thakur <1183441+abhishekkrthakur@users.noreply.github.com>
Date: Sun, 28 Jul 2024 14:19:58 +0200
Subject: [PATCH] Update Dockerfile

---
 Dockerfile | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/Dockerfile b/Dockerfile
index e9e5d3d808..4c2e8c6192 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -72,11 +72,10 @@ RUN conda install pytorch==2.3.0 torchvision==0.18.0 torchaudio==2.3.0 pytorch-c
     conda install -c "nvidia/label/cuda-12.1.1" cuda-nvcc && conda clean -ya && \
     conda install xformers -c xformers && conda clean -ya
 
-RUN pip install torch_xla~=2.3.0
-
 COPY --chown=1000:1000 . /app/
 
 RUN pip install -e . && \
+    pip install torch_xla~=2.3.0 && \
     python -m nltk.downloader punkt && \
     pip install -U ninja && \
     pip install -U flash-attn --no-build-isolation && \