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'] = [