From 8a34587e103beafc6cbb7e53afafb4c254a9e396 Mon Sep 17 00:00:00 2001
From: vfdev <vfdev.5@gmail.com>
Date: Thu, 7 Nov 2024 13:44:10 +0100
Subject: [PATCH] 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 | 2 +-
 docker/build.sh                             | 2 +-
 docker/docker.cfg                           | 2 +-
 docker/test_image.py                        | 5 +++++
 4 files changed, 8 insertions(+), 3 deletions(-)

diff --git a/.github/workflows/pytorch-version-tests.yml b/.github/workflows/pytorch-version-tests.yml
index 0ea40866aba..e14d46f2761 100644
--- a/.github/workflows/pytorch-version-tests.yml
+++ b/.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
diff --git a/docker/build.sh b/docker/build.sh
index f3b460a3943..b6a8495137d 100644
--- a/docker/build.sh
+++ b/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}
 
diff --git a/docker/docker.cfg b/docker/docker.cfg
index 6b976933347..7537dec68a6 100644
--- a/docker/docker.cfg
+++ b/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
diff --git a/docker/test_image.py b/docker/test_image.py
index ebb706c303f..77e1790d9c4 100644
--- 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}"