Resources for using PyTorch with Amazon SageMaker AI

You can use Amazon SageMaker AI to train and deploy a model using custom PyTorch code. The SageMaker AI Python SDK PyTorch estimators and models and the SageMaker AI open-source PyTorch container make writing a PyTorch script and running it in SageMaker AI easier. The following section provides reference material you can use to learn how to use PyTorch with SageMaker AI.

What do you want to do?

I want to train a custom PyTorch model in SageMaker AI.

For a sample Jupyter notebook, see the PyTorch example notebook in the Amazon SageMaker AI Examples GitHub repository.

For documentation, see Train a Model with PyTorch.

I have a PyTorch model that I trained in SageMaker AI, and I want to deploy it to a hosted endpoint.

For more information, see Deploy PyTorch models.

I have a PyTorch model that I trained outside of SageMaker AI, and I want to deploy it to a SageMaker AI endpoint

For more information, see Deploy your own PyTorch model.

I want to see the API documentation for Amazon SageMaker Python SDK PyTorch classes.

For more information, see PyTorch Classes.

I want to find the SageMaker AI PyTorch container repository.

For more information, see SageMaker AI PyTorch Container GitHub repository.

I want to find information about PyTorch versions supported by AWS Deep Learning Containers.

For more information, see Available Deep Learning Container Images.

For general information about writing PyTorch training scripts and using PyTorch estimators and models with SageMaker AI, see Using PyTorch with the SageMaker Python SDK.