pytorch_prodigy_optimizer_init_with_weight_decay.py

1import torch
2from prodigyopt import Prodigy
3
4# Define a simple model
5model = torch.nn.Linear(10, 1)
6
7# Initialize the Prodigy optimizer
8# Note: Prodigy typically requires a learning rate set to 1.0 
9# as it automatically estimates the optimal step size.
10optimizer = Prodigy(model.parameters(), lr=1.0, weight_decay=1e-4)
11
12# Example training step
13inputs = torch.randn(16, 10)
14targets = torch.randn(16, 1)
15
16optimizer.zero_grad()
17outputs = model(inputs)
18loss = torch.nn.functional.mse_loss(outputs, targets)
19loss.backward()
20optimizer.step()