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.

1. 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.

2. 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.

3. 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.

4. 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.

5. 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.

6. 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

7. Remove the smoke script when the AMP pattern has been copied into the project.

   $ rm amp_train_step.py