add torch version to the whl name (#24)

bytedance · Jul 22, 2024 · 532ad84 · 532ad84
1 parent 937bc18
commit 532ad84
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/gen_version.py b/gen_version.py
@@ -22,6 +22,7 @@
 import shutil
 import re
 from typing import Optional, Tuple
+import torch
 
 CUR_DIR = os.path.dirname(os.path.realpath(__file__))
 
@@ -83,7 +84,9 @@ def get_flux_version(version_txt, *, dev=False):
         version = f.readline()
         version = version.strip()
     cuda_version_major, cuda_version_minor = cuda_version()
-    version = version + f"+cu{cuda_version_major}{cuda_version_minor}"
+    torch_version_splits = torch.__version__.split(".")
+    torch_version = f"{torch_version_splits[0]}.{torch_version_splits[1]}"
+    version = version + f"+cu{cuda_version_major}{cuda_version_minor}" + f"torch{torch_version}"
     if dev:
         commit_id = get_git_commit(CUR_DIR)