Hello! I have a neural network like this (this is a simplified version):
class SimpleNet(nn.Module):
def __init__(self):
super().__init__()
self.linear1 = nn.Linear(2, 1, bias=False)
self.linear2 = nn.Linear(1, 2, bias=False)
def forward(self, x):
z = self.linear1(x)
y_pred = self.linear2(z)
return y_pred, z
and the loss is just MSE. After I call loss.backwards(), I want to get the value of the gradient of y_pred with respect to z. How can I do this? Thank you!