From 532ad8400f37a7cd8d690fd8288b18ecf556579f Mon Sep 17 00:00:00 2001
From: Ningxin Zheng <49771382+zheng-ningxin@users.noreply.github.com>
Date: Mon, 22 Jul 2024 15:54:37 +0800
Subject: [PATCH] add torch version to the whl name (#24)

---
 gen_version.py | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/gen_version.py b/gen_version.py
index c551d7863c..7091a24770 100644
--- 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)