I have a very basic model and i want to to print the loss and backward pass it. How can I use the forward function if it has more arguments than one?
Thank you
class Net(nn.Module):
def __init__(self,inp,hid,out):
super(Net,self).__init__()
self.layers = nn.ModuleDict(
{'lin1': nn.Linear(inp,hid),
'lin2': nn.Linear(hid,out)})
self.activations = nn.ModuleDict(
[['act1', nn.Sigmoid()],
['act2', nn.Sigmoid()]])
def forward(self, x, lay, activs):
x = self.layers[lay](x)
x = self.activations[activs](x)
return x
x = torch.randn(5, 4)
y = torch.randn(5, 3)
model = Net(4,3,3)
optimizer = torch.optim.SGD(model.parameters(),lr=0.0001)
loss_fn = nn.MSELoss(reduction='sum')
for t in range(500):
y_pred = model(x)
loss = loss_fn(y_pred, y)
print(t, loss.item())
optimizer.zero_grad()
loss.backward()
optimizer.step()