pytorch_ranger_optimizer_quickstart_simple_model_training.py

1import torch
2import torch.nn as nn
3from pytorch_ranger import Ranger
4
5# Define a simple model
6model = nn.Sequential(
7    nn.Linear(10, 50),
8    nn.ReLU(),
9    nn.Linear(50, 1)
10)
11
12# Initialize the Ranger optimizer
13# Ranger combines RAdam (Rectified Adam) and Lookahead
14optimizer = Ranger(model.parameters(), lr=1e-3)
15
16# Sample training step
17inputs = torch.randn(16, 10)
18targets = torch.randn(16, 1)
19
20optimizer.zero_grad()
21outputs = model(inputs)
22loss = nn.MSELoss()(outputs, targets)
23loss.backward()
24
25# Perform the optimization step
26optimizer.step()