update pytorch_local_mode_cifar10.ipynb to 1.7.1

sagemaker-python-sdk/pytorch_cnn_cifar10/pytorch_local_mode_cifar10.ipynb

@@ -199,7 +199,7 @@
     "\n",
     "cifar10_estimator = PyTorch(entry_point='source/cifar10.py',\n",
     "                            role=role,\n",
-    "                            framework_version='1.4.0',\n",
+    "                            framework_version='1.7.1',\n",
     "                            train_instance_count=1,\n",
     "                            train_instance_type=instance_type)\n",
     "\n",

sagemaker-python-sdk/pytorch_cnn_cifar10/source/cifar10.py

@@ -1,6 +1,5 @@
 import argparse
 import logging
-import sagemaker_containers
 
 import os
 
@@ -16,6 +15,11 @@
 import torchvision.transforms as transforms
 import torch.nn.functional as F
 
+try:
+    from sagemaker_inference import environment
+except:
+    from sagemaker_training import environment
+
 logger = logging.getLogger(__name__)
 logger.setLevel(logging.DEBUG)
 
@@ -150,7 +154,7 @@ def model_fn(model_dir):
     parser.add_argument('--momentum', type=float, default=0.9, metavar='M', help='momentum (default: 0.9)')
     parser.add_argument('--dist_backend', type=str, default='gloo', help='distributed backend (default: gloo)')
 
-    env = sagemaker_containers.training_env()
+    env = environment.Environment()
     parser.add_argument('--hosts', type=list, default=env.hosts)
     parser.add_argument('--current-host', type=str, default=env.current_host)
     parser.add_argument('--model-dir', type=str, default=env.model_dir)