How to set the Adam optimizer in PyTorch

Gradient-based training needs an optimizer to turn accumulated gradients into parameter updates. In PyTorch, Adam adapts each parameter's update using running first- and second-moment estimates, making it a common starting point when a fixed SGD step is too sensitive to gradient scale.

An Adam instance owns the parameters passed to its constructor and creates optimizer state for them when step() first sees gradients. Construct it from the same model parameters used by the forward pass; passing a different model's parameters leaves the active model unchanged.

The default learning rate is 1e-3; the small deterministic regression example uses 0.05 only to make one update easy to see. The final run prints the optimizer type and parameter-group learning rate, then checks state_step=1, a changed weight, and lower loss after the first update.

1. Create adam_step.py with deterministic data, a linear model, and its loss function.

   adam_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.Linear(3, 1)
   loss_fn = nn.MSELoss()

2. Append the Adam optimizer and initial weight snapshot below loss_fn.

   optimizer = torch.optim.Adam(model.parameters(), lr=0.05)
   weight_before = model.weight.detach().clone()

   Adam defaults to lr=1e-3. The higher value makes this one-step smoke run visible; select the learning rate from real training behavior instead of copying it into a project.
   Related: How to set a learning rate scheduler in PyTorch

3. Append one Adam optimization step below weight_before.

   optimizer.zero_grad(set_to_none=True)
   predictions = model(features)
   loss_before = loss_fn(predictions, targets)
   loss_before.backward()
   optimizer.step()

   zero_grad(set_to_none=True) prevents gradients from accumulating across iterations, while backward() fills gradients before step() updates the registered parameters.
   Related: How to zero gradients in PyTorch

4. Append the post-step state checks below optimizer.step().

   with torch.no_grad():
       loss_after = loss_fn(model(features), targets)
    
   state = optimizer.state[model.weight]
   weight_changed = not torch.equal(weight_before, model.weight.detach())
    
   print(f"optimizer={optimizer.__class__.__name__}")
   print(f"learning_rate={optimizer.param_groups[0]['lr']}")
   print(f"state_step={int(state['step'].item())}")
   print(f"weight_changed={weight_changed}")
   print(f"loss_before={loss_before.item():.6f}")
   print(f"loss_after={loss_after.item():.6f}")
   print(f"loss_decreased={loss_after.item() < loss_before.item()}")

5. Compare the completed adam_step.py source with the consolidated file.

   adam_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.Linear(3, 1)
   loss_fn = nn.MSELoss()
    
   optimizer = torch.optim.Adam(model.parameters(), lr=0.05)
   weight_before = model.weight.detach().clone()
    
   optimizer.zero_grad(set_to_none=True)
   predictions = model(features)
   loss_before = loss_fn(predictions, targets)
   loss_before.backward()
   optimizer.step()
    
   with torch.no_grad():
       loss_after = loss_fn(model(features), targets)
    
   state = optimizer.state[model.weight]
   weight_changed = not torch.equal(weight_before, model.weight.detach())
    
   print(f"optimizer={optimizer.__class__.__name__}")
   print(f"learning_rate={optimizer.param_groups[0]['lr']}")
   print(f"state_step={int(state['step'].item())}")
   print(f"weight_changed={weight_changed}")
   print(f"loss_before={loss_before.item():.6f}")
   print(f"loss_after={loss_after.item():.6f}")
   print(f"loss_decreased={loss_after.item() < loss_before.item()}")

6. Confirm the first Adam update by running adam_step.py.

   $ python adam_step.py
   optimizer=Adam
   learning_rate=0.05
   state_step=1
   weight_changed=True
   loss_before=0.179021
   loss_after=0.131637
   loss_decreased=True

   optimizer=Adam and learning_rate=0.05 identify the configured optimizer. state_step=1 proves that Adam initialized state for the model weight, while weight_changed=True and loss_decreased=True show that the same step changed the model and reduced this smoke loss.