feature: add inferentia pytorch inference container config (#1915)

* add inferentia pytorch config

* add test

* fix styling

Co-authored-by: Ashish Gupta <[email protected]>

Author: ashishgupta023

src/sagemaker/image_uri_config/inferentia-pytorch.json

@@ -0,0 +1,14 @@
+{
+    "processors": ["inf"],
+    "scope": ["inference"],
+    "versions": {
+        "1.5.1": {
+            "py_versions": ["py3"],
+            "registries": {
+                "us-east-1": "785573368785",
+                "us-west-2": "301217895009"
+            },
+            "repository": "sagemaker-neo-pytorch"
+        }
+    }
+}

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

@@ -118,6 +118,10 @@ 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(
         "sagemaker-{}".format(framework),