Optimizing regression weights for NN outputs with PyTorch
My solution is to update the variable manually and creating a new graph for autograd. Works perfectly well now!
my_nn.train()
for epoch in range(300):
    # zero the parameter gradients
    optimizer.zero_grad()
    my_nn.zero_grad()

    Y = my_nn(X)
    loss_output = criterion(Y, y)
    w = torch.zeros_like(target)
    w.requires_grad = True
    for nest_epoch in range(50):
        y_w = torch.matmul(X, w)
        loss_w = criterion(y_w, Y)
        grad = torch.autograd.grad(loss_w, w, create_graph=True)
        w = w - 0.1 * grad[0]
    loss_regression = criterion(w, target)
    loss = loss_output + loss_regression
    loss.backward()
    optimizer.step()