In my use case it’s
    def forward(self, input):
        out1 = self.sub_model2(self.sub_model1(input))
        out2 = self.sub_model3(out1)
        return out1, out2

That’s the tricky part. That’s why I want to zero out the gradient of sub_mode2 after calling backward on loss1 (loss1 is computed on out1).