diff --git a/tests/conftest.py b/tests/conftest.py
index 2c8dc2689f..7557c87fbe 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -254,6 +254,8 @@ 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.6"):
+        return "py312"
     if Version(pytorch_training_version) >= Version("2.3"):
         return "py311"
     elif Version(pytorch_training_version) >= Version("2.0"):
@@ -270,7 +272,9 @@ def pytorch_training_py_version(pytorch_training_version, request):
 
 @pytest.fixture(scope="module", params=["py2", "py3"])
 def pytorch_inference_py_version(pytorch_inference_version, request):
-    if Version(pytorch_inference_version) >= Version("2.3"):
+    if Version(pytorch_inference_version) >= Version("2.6"):
+        return "py312"
+    elif Version(pytorch_inference_version) >= Version("2.3"):
         return "py311"
     elif Version(pytorch_inference_version) >= Version("2.0"):
         return "py310"