fix: add pytorch 1.8.1 for huggingface

Author: Tabassum

src/sagemaker/image_uris.py

@@ -17,12 +17,9 @@
 import logging
 import os
 import re
-import pdb
 
 from sagemaker import utils
 from sagemaker.spark import defaults
-from sagemaker.spark import defaults
-
 
 logger = logging.getLogger(__name__)
 
@@ -106,7 +103,7 @@ def retrieve(
         instance_type, config.get("processors") or version_config.get("processors")
     )
     #if container version is available in .json file, utilize that
-    if "container_version" in version_config.keys():
+    if version_config.get("container_version"):
         container_version = version_config['container_version'][processor]
 
     if framework == HUGGING_FACE_FRAMEWORK:

tests/conftest.py

@@ -400,7 +400,7 @@ def _huggingface_base_fm_version(huggingface_version, base_fw, fixture_prefix):
             if len(original_version.split(".")) == 2:
                 base_fw_version = ".".join(base_fw_version.split(".")[:-1])
             versions.append(base_fw_version)
-    return versions
+    return sorted(versions,reverse=True)
 
 
 def _generate_huggingface_base_fw_latest_versions(