fix: flash attn build inside nvcr dockerfile (#646)

Signed-off-by: Dushyant Behl <dushyantbehl@in.ibm.com>

Author: dushyantbehl

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 \

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",