Update conftest.py

Author: eitansela

tests/conftest.py

@@ -173,6 +173,8 @@ def mxnet_eia_latest_py_version():
 def pytorch_training_py_version(pytorch_training_version, request):
     if Version(pytorch_training_version) < Version("1.5.0"):
         return request.param
+    elif Version(pytorch_training_version) == Version("1.7.1"):
+        return "py36"
     else:
         return "py3"
 
@@ -181,6 +183,8 @@ def pytorch_training_py_version(pytorch_training_version, request):
 def pytorch_inference_py_version(pytorch_inference_version, request):
     if Version(pytorch_inference_version) < Version("1.4.0"):
         return request.param
+    elif Version(pytorch_inference_version) == Version("1.7.1"):
+        return "py36"
     else:
         return "py3"