SageMaker Local Mode with PyTorch. Please run after data setup in https://github.com/aws-samples/amazon-sagemaker-examples-jp/blob/master/hpo_pytorch_mnist/pytorch_mnist.ipynb

pytorch_local.py

import os
import sagemaker

sagemaker_session = sagemaker.session.Session()

bucket = sagemaker_session.default_bucket()
prefix = 'sagemaker/DEMO-pytorch-mnist'

# set appropriate IAM Role
role = 'AmazonSageMaker-ExecutionRole-20210101T000000' 

from sagemaker.pytorch import PyTorch

estimator = PyTorch(entry_point="mnist.py",
                    role=role,
                    framework_version='1.6.0',
                    py_version='py3',
                    instance_count=1,
                    instance_type='local',
                    # instance_type='ml.p3.2xlarge',
                    hyperparameters={
                        'batch-size':128,
                        'lr': 0.01,
                        'epochs': 1,
                        'backend': 'gloo'
                    })
                    
estimator.fit({'training': os.path.join('s3://', bucket, prefix)})