Skip to content

Commit 532ad84

Browse files
add torch version to the whl name (#24)
1 parent 937bc18 commit 532ad84

File tree

1 file changed

+4
-1
lines changed

1 file changed

+4
-1
lines changed

gen_version.py

+4-1
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@
2222
import shutil
2323
import re
2424
from typing import Optional, Tuple
25+
import torch
2526

2627
CUR_DIR = os.path.dirname(os.path.realpath(__file__))
2728

@@ -83,7 +84,9 @@ def get_flux_version(version_txt, *, dev=False):
8384
version = f.readline()
8485
version = version.strip()
8586
cuda_version_major, cuda_version_minor = cuda_version()
86-
version = version + f"+cu{cuda_version_major}{cuda_version_minor}"
87+
torch_version_splits = torch.__version__.split(".")
88+
torch_version = f"{torch_version_splits[0]}.{torch_version_splits[1]}"
89+
version = version + f"+cu{cuda_version_major}{cuda_version_minor}" + f"torch{torch_version}"
8790
if dev:
8891
commit_id = get_git_commit(CUR_DIR)
8992

0 commit comments

Comments
 (0)