fix: add pytorch 1.8.1 for huggingface

Tabassum · Tabassum · commit e95e0ced5a9d · 2021-09-15T14:08:07.000-07:00
diff --git a/src/sagemaker/image_uris.py b/src/sagemaker/image_uris.py
@@ -122,7 +122,6 @@ def retrieve(
         container_version,
     )
 
-
     if _should_auto_select_container_version(instance_type, distribution):
         container_versions = {
             "tensorflow-2.3-gpu-py37": "cu110-ubuntu18.04-v3",
@@ -136,8 +135,7 @@ def retrieve(
             "pytorch-1.6-gpu-py36": "cu110-ubuntu18.04-v3",
             "pytorch-1.6.0-gpu-py36": "cu110-ubuntu18.04",
             "pytorch-1.6-gpu-py3": "cu110-ubuntu18.04-v3",
-            "pytorch-1.6.0-gpu-py3": "cu110-ubuntu18.04",
-            "pytorch-1.8.1-gpu-py3": "cu111-ubuntu18.04"
+            "pytorch-1.6.0-gpu-py3": "cu110-ubuntu18.04"
         }
 
 

Original file line number	Diff line number	Diff line change
`@@ -122,7 +122,6 @@ def retrieve(`
`122`	`122`	`container_version,`
`123`	`123`	`)`
`124`	`124`
`125`		`-`
`126`	`125`	`if _should_auto_select_container_version(instance_type, distribution):`
`127`	`126`	`container_versions = {`
`128`	`127`	`"tensorflow-2.3-gpu-py37": "cu110-ubuntu18.04-v3",`
`@@ -136,8 +135,7 @@ def retrieve(`
`136`	`135`	`"pytorch-1.6-gpu-py36": "cu110-ubuntu18.04-v3",`
`137`	`136`	`"pytorch-1.6.0-gpu-py36": "cu110-ubuntu18.04",`
`138`	`137`	`"pytorch-1.6-gpu-py3": "cu110-ubuntu18.04-v3",`
`139`		`- "pytorch-1.6.0-gpu-py3": "cu110-ubuntu18.04",`
`140`		`- "pytorch-1.8.1-gpu-py3": "cu111-ubuntu18.04"`
	`138`	`+ "pytorch-1.6.0-gpu-py3": "cu110-ubuntu18.04"`
`141`	`139`	`}`
`142`	`140`
`143`	`141`