fix: Fix ci unit-tests (aws#4713)

jiapinw · Jun 25, 2024 · 98e5acc · 98e5acc
1 parent bf2ebd2
commit 98e5acc
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/tests/conftest.py b/tests/conftest.py
@@ -253,7 +253,9 @@ def mxnet_eia_latest_py_version():
 
 @pytest.fixture(scope="module", params=["py2", "py3"])
 def pytorch_training_py_version(pytorch_training_version, request):
-    if Version(pytorch_training_version) >= Version("2.0"):
+    if Version(pytorch_training_version) >= Version("2.3"):
+        return "py311"
+    elif Version(pytorch_training_version) >= Version("2.0"):
         return "py310"
     elif Version(pytorch_training_version) >= Version("1.13"):
         return "py39"