How to enable mixed precision in PyTorch

Mixed precision in PyTorch lets a training loop run selected operations in a lower-precision dtype while model weights, gradients, and numerically sensitive work stay in safer precision. It is useful after the full-precision loop already trains correctly and the next target is lower memory use or better accelerator throughput.

Use the current torch.amp API for new code instead of the deprecated torch.cuda.amp or torch.cpu.amp namespaces. The autocast context belongs around the forward pass and loss calculation, while GradScaler belongs around the backward pass and optimizer step when the selected training dtype is float16.

A CPU smoke run can prove the same control flow with bfloat16 autocast when no CUDA GPU is available. Keep model weights in default precision, move inputs and targets to the selected device, and confirm that the optimizer updates parameters before applying the pattern to a larger training job.

Steps to enable PyTorch mixed precision:

Create a smoke script that runs one optimizer step with torch.amp.

amp_train_step.py

import torch
from torch import nn
 
 
torch.manual_seed(7)
 
device_type = "cuda" if torch.cuda.is_available() else "cpu"
device = torch.device(device_type)
amp_dtype = torch.float16 if device_type == "cuda" else torch.bfloat16
 
if device_type == "cpu":
    torch.backends.mkldnn.enabled = False
 
if not torch.amp.autocast_mode.is_autocast_available(device_type):
    raise SystemExit(f"autocast is not available for {device_type}")
 
model = nn.Sequential(
    nn.Linear(4, 8),
    nn.ReLU(),
    nn.Linear(8, 1),
).to(device)
optimizer = torch.optim.SGD(model.parameters(), lr=0.05)
loss_fn = nn.MSELoss()
 
inputs = torch.randn(16, 4, device=device)
target = torch.randn(16, 1, device=device)
scaler = torch.amp.GradScaler(device_type, enabled=(device_type == "cuda"))
 
optimizer.zero_grad(set_to_none=True)
before = model[0].weight.detach().clone()
 
with torch.amp.autocast(device_type=device_type, dtype=amp_dtype):
    prediction = model(inputs)
    loss = loss_fn(prediction, target)
 
scaler.scale(loss).backward()
grad_dtype = model[0].weight.grad.dtype
scaler.step(optimizer)
scaler.update()
 
weight_changed = not torch.equal(before, model[0].weight.detach())
 
print(f"torch_version={torch.__version__}")
print(f"device_type={device_type}")
print(f"autocast_available={torch.amp.autocast_mode.is_autocast_available(device_type)}")
print(f"amp_dtype={amp_dtype}")
print(f"scaler_enabled={scaler.is_enabled()}")
print(f"prediction_dtype={prediction.dtype}")
print(f"loss_dtype={loss.dtype}")
print(f"grad_dtype={grad_dtype}")
print(f"weight_changed={weight_changed}")
print(f"loss_value={loss.detach().float().item():.6f}")
print("step_complete=True")

The CPU branch uses bfloat16 autocast so the smoke script can run on a machine without a CUDA GPU. On CUDA, the same script uses float16 autocast and enables GradScaler.

Run the mixed-precision smoke script.
```
$ python amp_train_step.py
torch_version=2.12.1+cpu
device_type=cpu
autocast_available=True
amp_dtype=torch.bfloat16
scaler_enabled=False
prediction_dtype=torch.bfloat16
loss_dtype=torch.float32
grad_dtype=torch.float32
weight_changed=True
loss_value=0.648133
step_complete=True
```
prediction_dtype=torch.bfloat16 confirms that autocast selected a lower-precision output in the CPU smoke run. grad_dtype=torch.float32 and weight_changed=True confirm gradients reached the default-precision parameters and the optimizer updated them.

Add the device, dtype, and scaler setup to the real training script.

device_type = "cuda" if torch.cuda.is_available() else "cpu"
device = torch.device(device_type)
amp_dtype = torch.float16 if device_type == "cuda" else torch.bfloat16
scaler = torch.amp.GradScaler(device_type, enabled=(device_type == "cuda"))
 
model = model.to(device)

For ROCm builds, PyTorch exposes AMD GPUs through the torch.cuda device API, so the same cuda device type is used when torch.cuda.is_available() returns True.

Keep the forward pass and loss calculation inside autocast.
```
optimizer.zero_grad(set_to_none=True)
 
with torch.amp.autocast(device_type=device_type, dtype=amp_dtype):
    prediction = model(inputs.to(device))
    loss = loss_fn(prediction, target.to(device))
```
Do not call backward() inside the autocast context. Backward operations should use the dtype selected for the corresponding forward operations, and GradScaler needs to handle the loss outside the context.
Run the scaled backward pass and optimizer update outside autocast.
```
scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()
```
When scaler.is_enabled() is False, these calls behave like normal backward and optimizer calls. Keeping the same call shape lets CPU and CUDA smoke paths share one training-step function.

Check the first AMP step in the real training script.

first_parameter = next(model.parameters())
before = first_parameter.detach().clone()
 
optimizer.zero_grad(set_to_none=True)
with torch.amp.autocast(device_type=device_type, dtype=amp_dtype):
    prediction = model(inputs.to(device))
    loss = loss_fn(prediction, target.to(device))
 
scaler.scale(loss).backward()
grad_dtype = first_parameter.grad.dtype
scaler.step(optimizer)
scaler.update()
 
weight_changed = not torch.equal(before, first_parameter.detach())
 
print(f"prediction_dtype={prediction.dtype}")
print(f"grad_dtype={grad_dtype}")
print(f"weight_changed={weight_changed}")

If the loss, gradients, or model outputs become NaN or inf after enabling AMP, return to full precision and debug the failing tensor before changing the scaler or learning rate.
Related: How to debug NaN loss in PyTorch

Remove the smoke script when the AMP pattern has been copied into the project.
```
$ rm amp_train_step.py
```