Backport changes auto-install CUDA packages (#10402)

Oneflow-Inc · Jan 10, 2024 · 3bf53a0 · 3bf53a0
1 parent 9754f6d
commit 3bf53a0
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/python/setup.py b/python/setup.py
@@ -66,6 +66,10 @@ def get_version():
     REQUIRED_PACKAGES.append("nvidia-cudnn-cu11")
     REQUIRED_PACKAGES.append("nvidia-cublas-cu11")
 
+if "cu12" in ONEFLOW_VERSION:
+    REQUIRED_PACKAGES.append("nvidia-cudnn-cu12")
+    REQUIRED_PACKAGES.append("nvidia-cublas-cu12")
+
 # if python version < 3.7.x, than need pip install dataclasses
 if sys.version_info.minor < 7:
     REQUIRED_PACKAGES.append("dataclasses")