How to set the Adam optimizer in PyTorch

Optimizer choice controls how a PyTorch training loop turns gradients into parameter updates. Adam keeps running averages of gradients and squared gradients, so it is a common starting point for small models, noisy losses, and custom loops that need adaptive learning rates instead of plain SGD.

torch.optim.Adam receives model parameters and stores per-parameter optimizer state after the first successful step() call. Build it after the model has been created and moved to the target device, then run it in the usual training order of clearing old gradients, calculating loss, backpropagating, and stepping the optimizer.

A CPU smoke run with a tiny linear model makes the optimizer state visible without accelerator-specific setup. A changed weight, state_step=1, non-zero exp_avg_norm, and lower loss confirm that Adam updated the model and initialized its first-moment state.

1. Create an Adam optimizer smoke script.

   adam_optimizer_step.py

   import torch
   from torch import nn
    
    
   torch.manual_seed(7)
    
   features = torch.tensor(
       [
           [-1.0, 0.5, 0.25],
           [0.0, -0.25, 0.75],
           [0.5, 0.75, -0.5],
           [1.0, -0.5, -0.25],
       ],
       dtype=torch.float32,
   )
   targets = features @ torch.tensor([[0.6], [-0.4], [0.2]]) + 0.1
    
   model = nn.Sequential(nn.Linear(3, 1))
   loss_fn = nn.MSELoss()
   optimizer = torch.optim.Adam(model.parameters(), lr=0.05)
    
   first_weight = model[0].weight
   weight_before = first_weight.detach().clone()
    
   optimizer.zero_grad(set_to_none=True)
   prediction = model(features)
   loss_before = loss_fn(prediction, targets)
   loss_before.backward()
    
   gradient_norm = torch.linalg.vector_norm(first_weight.grad).item()
   optimizer.step()
    
   with torch.no_grad():
       loss_after = loss_fn(model(features), targets)
    
   state = optimizer.state[first_weight]
   weight_changed = not torch.equal(weight_before, first_weight.detach())
    
   print(f"torch_version={torch.__version__}")
   print(f"optimizer={optimizer.__class__.__name__}")
   print(f"learning_rate={optimizer.param_groups[0]['lr']}")
   print(f"loss_before={loss_before.item():.6f}")
   print(f"gradient_norm={gradient_norm:.6f}")
   print(f"state_step={int(state['step'].item())}")
   print(f"exp_avg_norm={torch.linalg.vector_norm(state['exp_avg']).item():.6f}")
   print(f"weight_changed={weight_changed}")
   print(f"loss_after={loss_after.item():.6f}")
   print(f"loss_decreased={loss_after.item() < loss_before.item()}")

   lr=0.05 is intentionally high for a tiny smoke test. Use a project-specific learning rate, often near 1e-3, before tuning a real model.

2. Run the smoke script.

   $ python adam_optimizer_step.py
   torch_version=2.12.1+cpu
   optimizer=Adam
   learning_rate=0.05
   loss_before=0.179021
   gradient_norm=0.690800
   state_step=1
   exp_avg_norm=0.069080
   weight_changed=True
   loss_after=0.131637
   loss_decreased=True

   weight_changed=True confirms that optimizer.step() updated a model parameter. state_step=1 and exp_avg_norm confirm that Adam created optimizer state for the first tracked parameter.

3. Create the Adam optimizer after the model parameters exist.

   model = build_model().to(device)
   optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)

   Pass model.parameters() directly when every trainable parameter uses the same options. Use parameter groups only when different layers need different learning rates or weight decay values.

4. Clear old gradients at the start of each optimization step.

   optimizer.zero_grad(set_to_none=True)

   set_to_none=True lets parameters that receive no gradient stay at None, and torch.optim skips those parameters during the step.
   Related: How to zero gradients in PyTorch

5. Compute the loss and gradients.

   predictions = model(batch_features)
   loss = loss_fn(predictions, batch_targets)
   loss.backward()

6. Update the model parameters with Adam.

   optimizer.step()

   Call learning-rate schedulers after the optimizer update unless a specific scheduler documents another order.
   Related: How to set a learning rate scheduler in PyTorch

7. Check the first real optimizer step while wiring the loop.

   first_parameter = next(model.parameters())
   before = first_parameter.detach().clone()
    
   optimizer.zero_grad(set_to_none=True)
   predictions = model(batch_features)
   loss = loss_fn(predictions, batch_targets)
   loss.backward()
   optimizer.step()
    
   weight_changed = not torch.equal(before, first_parameter.detach())
   print(f"weight_changed={weight_changed}")

   If weight_changed is False, inspect whether the loss used model outputs, whether gradients are None, and whether the optimizer was constructed with the same parameters used in the forward pass.

8. Remove the temporary smoke script after the pattern is in the project.

   $ rm adam_optimizer_step.py