add manual torch install to the dockerfile

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

Author: dushyantbehl

build/Dockerfile

@@ -86,6 +86,7 @@ ENV CUDA_HOME="/usr/local/cuda" \
 ## CUDA Development ############################################################
 FROM cuda-base AS cuda-devel
 
+# Note:- check the torch installation below as that is hardcoded to cuda 12.1
 # Ref: https://developer.nvidia.com/nccl/nccl-legacy-downloads
 ENV NV_CUDA_CUDART_DEV_VERSION=12.1.55-1 \
     NV_NVML_DEV_VERSION=12.1.55-1 \
@@ -155,8 +156,10 @@ 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 --mount=type=cache,target=/home/${USER}/.cache/pip,uid=${USER_UID} \
+    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