Getting Started With Mixed Precision Using torch.amp
torch.amp, introduced in PyTorch 1.6, makes it easy to leverage mixed precision training using the float16 or bfloat16 dtypes. See this blog post, tutorial, and documentation for more details. Figure 4 shows an example of applying AMP with grad scaling to a network.
```console
import torch
Creates once at the beginning of training
scaler = torch.cuda.amp.GradScaler()
for data, label in data_iter:
   optimizer.zero_grad()
   # Casts operations to mixed precision
   with torch.amp.autocast(device_type=“cuda”, dtype=torch.float16):
      loss = model(data)
# Scales the loss, and calls backward()
   # to create scaled gradients
   scaler.scale(loss).backward()
# Unscales gradients and calls
   # or skips optimizer.step()