I am getting ValueError: optimizer got an empty parameter list
when I run this code. What am I missing? Tried using nn.Parameters
and still getting error on model parameters a
and b
. I just want to track and update these two scalars on this toy problem.
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import TensorDataset, DataLoader
class RosenbrockFunction2N(nn.Module):
def __init__(self):
super(RosenbrockFunction2N, self).__init__()
self.a = torch.rand(1, requires_grad=True)
self.b = torch.rand(1, requires_grad=True)
def forward(self, x, y):
return (self.a - x)**2 + self.b(y-x**2)**2
# Dataset
inputs = [[1., 3.], [5., 7.], [2., 4.], [3., 9.]]
targets = []
for i in inputs:
targets.append([(1-i[0]**2) + (i[1] - i[0]**2)**2])
inputs = torch.tensor(inputs)
targets = torch.tensor(targets)
train_ds = TensorDataset(inputs, targets)
train_dl = DataLoader(train_ds)
# Load model
model = RosenbrockFunction2N()
# Loss and Optimizer
loss_fn = nn.MSELoss()
optimizer = torch.optim.SGD(model.parameters(), lr=0.01, momentum=0.9)
for i in range(5):
for inputs, target in train_dl:
optimizer.zero_grad()
output = model(inputs[0][0], inputs[0][1])
loss = loss_fn(output, target)
loss.backward()
optimizer.step()