albert.cwkuo:

My goal is to accumulate gradient from both loss1 and loss2 to sub_model1 and accumulate gradient only from loss2 to sub_model2 and sub_model3.


With the above statement, should the forward function be something like below?
    def forward(self, input):
        out1 = self.sub_model1(input)
        out2 = self.sub_model3(self.sub_model2(out1))
        return out1, out2

With this code, out1.sum().backward() will only compute grads for sub_model1, and out2.sum().backward() will compute grads for all sub-models. And (out1 + out2).sum().backward() should meet the cited statement above.