Skip to content

Commit 9982f57

Browse files
HappyAmazonianmufiAmazon
authored andcommitted
feature: Add Neo image uri config for Pytorch 1.12 (aws#3507)
1 parent 20a4ade commit 9982f57

File tree

2 files changed

+37
-3
lines changed

2 files changed

+37
-3
lines changed

src/sagemaker/image_uri_config/neo-pytorch.json

+35-1
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,9 @@
1111
"1.7.0": "1.7",
1212
"1.7.1": "1.7",
1313
"1.8.0": "1.8",
14-
"1.8.1": "1.8"
14+
"1.8.1": "1.8",
15+
"1.12.0": "1.12",
16+
"1.12.1": "1.12"
1517
},
1618
"versions": {
1719
"1.4": {
@@ -173,6 +175,38 @@
173175
"us-west-2": "301217895009"
174176
},
175177
"repository": "sagemaker-inference-pytorch"
178+
},
179+
"1.12": {
180+
"py_versions": ["py3"],
181+
"registries": {
182+
"af-south-1": "774647643957",
183+
"ap-east-1": "110948597952",
184+
"ap-northeast-1": "941853720454",
185+
"ap-northeast-2": "151534178276",
186+
"ap-northeast-3": "925152966179",
187+
"ap-south-1": "763008648453",
188+
"ap-southeast-1": "324986816169",
189+
"ap-southeast-2": "355873309152",
190+
"ca-central-1": "464438896020",
191+
"cn-north-1": "472730292857",
192+
"cn-northwest-1": "474822919863",
193+
"eu-central-1": "746233611703",
194+
"eu-north-1": "601324751636",
195+
"eu-south-1": "966458181534",
196+
"eu-west-1": "802834080501",
197+
"eu-west-2": "205493899709",
198+
"eu-west-3": "254080097072",
199+
"me-south-1": "836785723513",
200+
"sa-east-1": "756306329178",
201+
"us-east-1": "785573368785",
202+
"us-east-2": "007439368137",
203+
"us-gov-west-1": "263933020539",
204+
"us-iso-east-1": "167761179201",
205+
"us-isob-east-1": "406031935815",
206+
"us-west-1": "710691900526",
207+
"us-west-2": "301217895009"
208+
},
209+
"repository": "sagemaker-inference-pytorch"
176210
}
177211
}
178212
}

tests/data/pytorch_neo/code/inference.py

+2-2
Original file line numberDiff line numberDiff line change
@@ -71,8 +71,8 @@ def model_fn(model_dir):
7171
logger.info("model_fn")
7272
neopytorch.config(model_dir=model_dir, neo_runtime=True)
7373
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
74-
# The compiled model is saved as "model.pth"
75-
model = torch.jit.load(os.path.join(model_dir, "model.pth"), map_location=device)
74+
# The compiled model is saved as "model.pth" or "model.pt"
75+
model = torch.jit.load(os.path.join(model_dir, "model.pt"), map_location=device)
7676

7777
# It is recommended to run warm-up inference during model load
7878
sample_input_path = os.path.join(model_dir, "sample_input.pkl")

0 commit comments

Comments
 (0)