Fix version compare in test_image

pytorch · Nov 7, 2024 · 62a080f · 62a080f
1 parent cd7cc11
commit 62a080f
Showing 1 changed file with 5 additions and 0 deletions.
diff --git a/docker/test_image.py b/docker/test_image.py
@@ -19,6 +19,11 @@ def check_package(package_name, expected_version=None):
     if expected_version is not None:
         assert hasattr(mod, "__version__"), f"Imported package {package_name} does not have __version__ attribute"
         version = mod.__version__
+        # Remove all +something from the version name: e.g torch 2.5.1+cu124
+        if "+" in version:
+            old_version = version
+            version = version.split("+")[0]
+            print(f"Transformed version: {old_version} -> {version}")
         assert (
             version == expected_version
         ), f"Version mismatch for package {package_name}: got {version} but expected {expected_version}"