Hey, since I am stuck with the error # RuntimeError: Function 'AddmmBackward0' returned nan values in its 0th output. and there seems to be no way of getting around it, I would like to set the backward graph by myself.
How would you do that?
You can write a custom autograd.Function defining your own backward as described here.