Address review comments

Signed-off-by: Thara Palanivel <130496890+tharapalanivel@users.noreply.github.com>

Author: tharapalanivel

build/Dockerfile

@@ -86,9 +86,12 @@ RUN python -m pip install packaging && \
     python -m pip install --upgrade pip && \
     python -m pip install torch && \
     python -m pip install wheel
+
+# TODO Move to installing wheel once we have proper releases setup instead of cloning the repo
 RUN git clone https://github.com/foundation-model-stack/fms-hf-tuning.git && \
     cd fms-hf-tuning && \
     python -m pip install -r requirements.txt && \
+    python -m pip install -r flashattn_requirements.txt && \
     python -m pip install -U datasets  && \
     python -m pip install /tmp/fms-hf-tuning