diff --git a/build/nvcr.Dockerfile b/build/nvcr.Dockerfile
index e277b7f25b..1ee94d5e2e 100644
--- a/build/nvcr.Dockerfile
+++ b/build/nvcr.Dockerfile
@@ -47,8 +47,9 @@ RUN pip install --upgrade --force-reinstall torch torchaudio torchvision --index
 # Install main package + flash attention
 COPY . ${SOURCE_DIR}
 RUN cd ${SOURCE_DIR}
-RUN pip install --no-cache-dir ${SOURCE_DIR} && \
-    pip install --no-cache-dir ${SOURCE_DIR}[flash-attn]
+
+RUN pip install --no-cache-dir ${SOURCE_DIR}
+RUN pip install --user --no-build-isolation ${SOURCE_DIR}[flash-attn]
 
 # Optional extras
 RUN if [[ "${ENABLE_FMS_ACCELERATION}" == "true" ]]; then \
diff --git a/pyproject.toml b/pyproject.toml
index 9d09a36b34..8f7f6b2de1 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -30,6 +30,7 @@ dependencies = [
 "accelerate>=1.9.0,<2.0.0",
 "transformers>=4.55.0,<=4.55.4",
 "torch>2.7.0,<2.9.0",
+"torchvision<0.24",
 "sentencepiece>=0.1.99,<0.3",
 "tokenizers<=0.22",
 "tqdm>=4.66.2,<5.0",