Updated pytorch version in CI and docker (#3298)

* Updated pytorch version in CI and docker

* Fix docker tag

* Fix docker tag value removing +cu124 like words from torch version

* Fix version compare in test_image

.github/workflows/pytorch-version-tests.yml

@@ -19,7 +19,7 @@ jobs:
         # will drop python version and related pytorch versions
         python-version: [3.8, 3.9, "3.10"]
         pytorch-version:
-          [2.3.1, 2.2.2, 2.1.2, 2.0.1, 1.13.1, 1.12.1, 1.10.0, 1.8.1]
+          [2.4.1, 2.3.1, 2.2.2, 2.0.1, 1.13.1, 1.12.1, 1.10.0, 1.8.1]
         exclude:
           # disabling python 3.9 support with PyTorch 1.7.1 and 1.8.1, to stop repeated pytorch-version test fail.
           # https://github.com/pytorch/ignite/issues/2383

docker/build.sh

@@ -80,7 +80,7 @@ echo "opt_build_args: ${opt_build_args}"
 
 retry "docker build --build-arg PTH_VERSION=${pth_version} ${opt_build_args} -t pytorchignite/${image_name}:latest -f Dockerfile.${image_name} ." "\nBuild failed: ${image_name}"
 if [ -z $image_tag ]; then
-    image_tag=`docker run --rm -i pytorchignite/${image_name}:latest python -c "import torch; import ignite; print(torch.__version__ + \"-\" + ignite.__version__, end=\"\")"`
+    image_tag=`docker run --rm -i pytorchignite/${image_name}:latest python -c "import torch; import ignite; print(torch.__version__.split('+')[0] + \"-\" + ignite.__version__, end=\"\")"`
 fi
 docker tag pytorchignite/${image_name}:latest pytorchignite/${image_name}:${image_tag}
 

docker/docker.cfg

@@ -1,4 +1,4 @@
 [DEFAULT]
-build_docker_image_pytorch_version = 2.4.0-cuda12.4-cudnn9
+build_docker_image_pytorch_version = 2.5.1-cuda12.4-cudnn9
 build_docker_image_hvd_version = v0.28.1
 build_docker_image_msdp_version = v0.14.0

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}"