add manual torch install to the dockerfile

Signed-off-by: Dushyant Behl <[email protected]>

Author: dushyantbehl

build/Dockerfile

@@ -155,8 +155,9 @@ ENV PIP_NO_BINARY=mamba-ssm,mamba_ssm
 # Install from the wheel
 RUN --mount=type=cache,target=/home/${USER}/.cache/pip,uid=${USER_UID} \
     python -m pip install --user wheel && \
-    python -m pip install --user "$(head bdist_name)" && \
-    python -m pip install --user "$(head bdist_name)[flash-attn]" && \
+    python -m pip install --user "$(head bdist_name)"
+
+RUN python -m pip install --user --no-build-isolation "$(head bdist_name)[flash-attn]" && \
     python -m pip install --user --no-build-isolation "$(head bdist_name)[mamba]"
 
 # fms_acceleration_peft = PEFT-training, e.g., 4bit QLoRA