Assume that I have model Model like this
layer1->layer2->layer3.
I feed a input in the Model, then I have
input -> layer1 -> output1 -> layer2 -> output2 -> layer3 -> output3
After I call output3.sum().backward()
, the default gradient computation is as follows
output3.grad -> layer3.grad -> output2.grad -> layer2.grad -> output1.grad -> layer1.grad
However, I want to schedule this gradient computation order like this
output3.grad -> output2.grad -> output1.grad -> layer1.grad->layer2.grad->layer3.grad
Is it any suggestion about how to achieve this?