Hi,
I am trying to implement a custom loss function that involves computing an inner optimization loop. I am using a torch.autograd.Function because I would like to manually compute the gradient.
The function looks something like this:
# Inside the forward method of the Function
output_theta = model(X)
theta_star = {}
theta_inital = deepcopy(model.state_dict())
for name, v in theta_inital.items():
    random_val = torch.randn(size=v.shape, device=device)
    theta_star[name] = torch.tensor(v + v.abs() * initial_std * random_val, requires_grad=True)

for _ in range(n_attack_steps):
    for name, v in model.named_parameters():
        v.copy_(theta_star[name])
    output_theta_star = model(X)
    step_loss = robustness_loss(torch.nn.functional.softmax(output_theta, dim=1).log(), output_theta_star)
    step_loss.backward()

The implementation is not done, but the problem that I now have is that I can’t get any gradients on theta_star.
Is it generally not possible to calculate gradients inside a Function?
I have the feeling that using a Function is not the way to go here, but I need to implement the gradient in a custom fashion.
As a reference, I want to implement the loss function of this paper: https://arxiv.org/pdf/2106.05009.pdf
What do I do?