Hey all,
I am having a problem with the autograd for different losses in diffrent layers, I want to implement a new algorithm, but first I did a check test to see if results are the same.
import torch 
import torch.nn as nn

device = 'cpu'

#define neural network
class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.layer1 = nn.Sequential(
            nn.Linear(1, 1, bias=False))
        self.layer2 = nn.Sequential(
            nn.Linear(1, 1, bias=False))
        
    def forward(self, x):
        out1 = self.layer1(x)
        out2 = self.layer2(out1)
        return out1,out2
    
model = Net().to(device)

#desired output
t = (torch.randn(1, 1)).to(device)

#input
inp = (torch.randn(1, 1)).to(device)

#run network
outputs = model(inp)

#loss from the ouput layer
loss = 0.5*torch.sum( (t - outputs[1])**2.0)

#calculate the gradients just to check
loss.backward(retain_graph=True)

#gradient of the second weight
grad1 = torch.autograd.grad(loss, model.layer2[0].weight, retain_graph=True)

#part of the chain rule to calculate the gradient of the first weight
t1 = torch.autograd.grad(-loss, outputs[0], create_graph=True, retain_graph=True)[0]

#so it will be same as backpropagation
t1 += outputs[0]

#new loss for the gradient
loss2 = 0.5*torch.sum( (t1 - outputs[0])**2.0)

#gradient for the first weight
grad2 = torch.autograd.grad(loss2, model.layer1[0].weight)

print(grad2[0])
print(model.layer1[0].weight.grad)

loss2 is the loss function for the hidden layer. Mathematically this should be the same as the usual backpropagation, the results should be equal, but they aren’t. I checked t1 and it is all right, but something is going on with grad2, can someone help me with this?