I have a question regarding training a model’s parameters when the forward function include a gradient wrt to the input. Instead of showing all the code (where I have trouble with the backward-function), I can boil it down to this small (and silly) example (and when I can solve this problem I think it will translate to my code):
import torch as th
import torch.nn as nn
import torch.nn.functional as F
def __init__(self): super().__init__() self.l1 = nn.Linear(5, 2) def forward(self, x): x_ = self.l1(x) x_ = (x_ ** 2).sum(1) return th.autograd.grad(x_, x, grad_outputs=th.ones_like(x_))
m = Model()
x = th.randn(10, 5).requires_grad_(True)
y = m(x)
loss = F.mse_loss(y, x)
params = m.parameters()
for p in params:
print(th.autograd.grad(loss, p, retain_graph=True))
Here I cannot compute the gradient in the for-loop. How can I solve this?
Thanks in advance!!