add test

Author: Ashish Gupta

tests/conftest.py

@@ -44,6 +44,7 @@
     "coach_tensorflow",
     "inferentia_mxnet",
     "inferentia_tensorflow",
+    "inferentia_pytorch",
     "mxnet",
     "neo_mxnet",
     "neo_pytorch",

tests/unit/sagemaker/image_uris/test_neo.py

@@ -117,6 +117,8 @@ def test_inferentia_mxnet(inferentia_mxnet_version):
 def test_inferentia_tensorflow(inferentia_tensorflow_version):
     _test_inferentia_framework_uris("tensorflow", inferentia_tensorflow_version)
 
+def test_inferentia_pytorch(inferentia_pytorch_version):
+    _test_inferentia_framework_uris("pytorch", inferentia_pytorch_version)
 
 def _expected_framework_uri(framework, version, region="us-west-2", processor="cpu"):
     return expected_uris.framework_uri(