Oh, thanks for pointing that out. I guess I had a misunderstanding of the gradients (actually its my very first model ). But for example if we are passing inputs through something like.
L = nn.Linear(somethingsomething)
x = L(x)
x = x^2
x = log(x)
shouldn’t gradients be calculated first with respect to x=log(x) and x=x^2, then the linear layer, rather than directly getting loss derivatives wrt whatever loss values outputted by the loss function? Thanks!