Fortunately, PyTorch provides methods and tools to perform AMP by changing just a few things in our original code.

In PyTorch, AMP relies on enabling a couple of flags, wrapping the training process with the torch.autocast object, and using a gradient scaler. The more complex case, which is related to implementing AMP on GPU, takes all these three parts, while the most simple scenario (CPU-based training) requires only the usage of torch.autocast.

Let’s start by covering the more complex scenario. So, follow me to the next section to learn how to activate this approach in our GPU-based code.

Activating AMP on GPU

To activate AMP on GPU, we need to make three modifications to our code:

1. Enable the CUDA and CuDNN backend flags.
2. Wrap the training loop with torch.autocast.
3. Use a gradient scaler.

Let’s take a closer look.

Enabling backend flags

As we learned in Chapter 4, Using Specialized Libraries, PyTorch relies on third...