ci: Fixes for pre-built wheels (#214)

* build: Allow NGC builds Signed-off-by: oliver könig <okoenig@nvidia.com> * reduce grid Signed-off-by: oliver könig <okoenig@nvidia.com> * update grid Signed-off-by: oliver könig <okoenig@nvidia.com> * fix Signed-off-by: oliver könig <okoenig@nvidia.com> * upgrade cuda action Signed-off-by: oliver könig <okoenig@nvidia.com> * remove test Signed-off-by: oliver könig <okoenig@nvidia.com> * py3.8 Signed-off-by: oliver könig <okoenig@nvidia.com> * fix Signed-off-by: oliver könig <okoenig@nvidia.com> * exclude Signed-off-by: oliver könig <okoenig@nvidia.com> * fix Signed-off-by: oliver könig <okoenig@nvidia.com> * torch-version Signed-off-by: oliver könig <okoenig@nvidia.com> * py3.8/torch2.1/cuda12.3 Signed-off-by: oliver könig <okoenig@nvidia.com> * Update publish.yml * fix grid Signed-off-by: oliver könig <okoenig@nvidia.com> * fix Signed-off-by: oliver könig <okoenig@nvidia.com> * cuda11.8 Signed-off-by: oliver könig <okoenig@nvidia.com> * no hopper for 118 Signed-off-by: oliver könig <okoenig@nvidia.com> * fix Signed-off-by: oliver könig <okoenig@nvidia.com> * fix Signed-off-by: oliver könig <okoenig@nvidia.com> --------- Signed-off-by: oliver könig <okoenig@nvidia.com>
2025-10-14 07:05:47 +02:00
parent f8f41145da
commit 93b3c28fa8
8 changed files with 337 additions and 34 deletions
--- a/setup.py
+++ b/setup.py
@@ -84,7 +84,11 @@ def get_platform():

 def get_wheel_url():
    torch_version = parse(torch.__version__)
-    torch_version = f'{torch_version.major}.{torch_version.minor}'
+    if os.environ.get("NVIDIA_PRODUCT_NAME", "") == "PyTorch":
+        torch_version = str(os.environ.get("NVIDIA_PYTORCH_VERSION"))
+    else:
+        torch_version = f'{torch_version.major}.{torch_version.minor}'
+
    python_version = f'cp{sys.version_info.major}{sys.version_info.minor}'
    platform_name = get_platform()
    deep_gemm_version = get_package_version()