From 99cea1460ebfa54b64e86a82ecae48fd88b6e8f9 Mon Sep 17 00:00:00 2001
From: Mihir Patel <mihir.v.patel7@gmail.com>
Date: Wed, 20 Sep 2023 21:13:58 -0400
Subject: [PATCH] Bump flash attn (#2551)

* bump flash attn

* onnx

* bump to 1.0.9

* pin

* diff

---------

Co-authored-by: Daniel King <43149077+dakinggg@users.noreply.github.com>
---
 docker/Dockerfile | 2 +-
 setup.py          | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/docker/Dockerfile b/docker/Dockerfile
index bf128a5025..dcc672d046 100644
--- a/docker/Dockerfile
+++ b/docker/Dockerfile
@@ -305,7 +305,7 @@ RUN if [[ -n "$CUDA_VERSION" ]] &&  [[ -z "${PYTORCH_NIGHTLY_URL}" ]]; then \
 RUN if [ -n "$CUDA_VERSION" ] ; then \
         pip${PYTHON_VERSION} install --upgrade --no-cache-dir ninja==1.11.1 && \
         pip${PYTHON_VERSION} install --upgrade --no-cache-dir --force-reinstall packaging==22.0 && \
-        pip${PYTHON_VERSION} install --no-cache-dir flash-attn==1.0.3.post0; \
+        pip${PYTHON_VERSION} install --no-cache-dir flash-attn==1.0.9; \
     fi
 
 ###############
diff --git a/setup.py b/setup.py
index a45e90a3da..e725683756 100644
--- a/setup.py
+++ b/setup.py
@@ -219,7 +219,7 @@ def package_files(prefix: str, directory: str, extension: str):
 
 extra_deps['onnx'] = [
     'onnx>=1.12.0,<2',
-    'onnxruntime>=1.12.1,<2',
+    'onnxruntime>=1.12.1,<1.16',
 ]
 
 extra_deps['mlflow'] = [