From 0d037e0ed5dc86e6c4a27d62e4e483d1c49ec86c Mon Sep 17 00:00:00 2001 From: tsai Date: Mon, 6 May 2024 16:30:26 +0800 Subject: [PATCH] update deps --- python/setup.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/python/setup.py b/python/setup.py index 93b600e50dd..4ce7eab7d99 100644 --- a/python/setup.py +++ b/python/setup.py @@ -13,6 +13,7 @@ See the License for the specific language governing permissions and limitations under the License. """ + from __future__ import absolute_import import argparse @@ -63,14 +64,14 @@ def get_version(): ONEFLOW_VERSION = get_version() if "cu11" in ONEFLOW_VERSION and "cu112" not in ONEFLOW_VERSION: - REQUIRED_PACKAGES.append("nvidia-cudnn-cu11") + REQUIRED_PACKAGES.append("nvidia-cudnn-cu11<9") REQUIRED_PACKAGES.append("nvidia-cublas-cu11") REQUIRED_PACKAGES.append("nvidia-nccl-cu11") REQUIRED_PACKAGES.append("nvidia-cusparse-cu11") REQUIRED_PACKAGES.append("nvidia-cufft-cu11") if "cu12" in ONEFLOW_VERSION: - REQUIRED_PACKAGES.append("nvidia-cudnn-cu12") + REQUIRED_PACKAGES.append("nvidia-cudnn-cu12<9") REQUIRED_PACKAGES.append("nvidia-cublas-cu12") REQUIRED_PACKAGES.append("nvidia-nccl-cu12") REQUIRED_PACKAGES.append("nvidia-cusparse-cu12")