albert.cwkuo:

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


In that case, calling backward once on loss1+loss2 might be sufficient? Is the following result what you want?
import torch
import torch.nn as nn

class MyModel(nn.Module):

    def __init__(self):
        super().__init__()
        with torch.no_grad():
            self.net1 = nn.Linear(1, 1)
            self.net1.weight.copy_(torch.ones(1, 1))
            self.net1.bias.copy_(torch.zeros(1))
            self.net2 = nn.Linear(1, 1)
            self.net2.weight.copy_(torch.ones(1, 1))
            self.net2.bias.copy_(torch.zeros(1))

    def forward(self, x):
        out1 = self.net1(x)
        out2 = self.net2(out1)
        return out1 + out2


print("==============")
model = MyModel()
model(torch.ones(1, 1)).sum().backward()
print("net1 grad is: ", model.net1.weight.grad)
print("net2 grad is: ", model.net2.weight.grad)

print("==============")
model = MyModel()
model.net1(torch.ones(1, 1)).sum().backward()
print("net1 grad is: ", model.net1.weight.grad)
print("net1 grad is: ", model.net2.weight.grad)

print("==============")
model = MyModel()
model.net2(torch.ones(1, 1)).sum().backward()
print("net1 grad is: ", model.net1.weight.grad)
print("net2 grad is: ", model.net2.weight.grad)

outputs are
==============
net1 grad is:  tensor([[2.]])
net2 grad is:  tensor([[1.]])
==============
net1 grad is:  tensor([[1.]])
net2 grad is:  None
==============
net1 grad is:  None
net2 grad is:  tensor([[1.]])