I’m trying out meta learning in PyTorch, and I’m having trouble propagating gradients through the parameters of an nn.Module. Here’s a simplified example that’s not working as I had expected:
import torch
from torch import nn

net = nn.Linear(2, 1)

x = torch.tensor([[1., 2.]])
y = torch.tensor([[2.]])

loss = (net(x) - y)**2
grads = torch.autograd.grad(loss, net.parameters(), create_graph=True)

net2 = nn.Linear(2, 1)
net2.weight.data = net2.weight.data - 0.01*grads[0]
net2.bias.data = net2.bias.data - 0.01*grads[1]

loss2 = (net2(x) - y)**2
loss2.backward()

print net2.weight.grad # This is populated as expected
print net.weight.grad # This is 'None', which is not what I'd like

How do I get gradients into net.weight?