Thanks for the fast reply, and I’m always happy to learn!


I can not exchange 1 by 1 the functions, as what is returned as a “gradient” is the concatenation of the gradient + identity. Thus if I plug in only one of my Functions, the graph will break since the shapes will not match.


Thanks, I’ve now changed everything to backward(self, grad_output) and unpack them from tuples (Updated the code in the post as well).


I should have mentioned that the error occurs immediately after the return statement of Gn_SquareLoss.backward so non of the other function’s backward methods get called.


To your question - if the grad is standard gradient is NxD of the loss, I append an identity on the 0th axis, thus now return (N+D)xD. Then in the backward of each other functions I unpack that two a DxD and NxD tensors via indexing as you said.


Hope this make it clear.