Thanks for your response!
But with your approach, when you do
loss = loss_A + loss_B
loss.backward()

both part A and B will calculate gradient for loss_A and loss_B, and add them together. And when you do
optimizer_B.step()

optimizer_B will update part B with loss_A. This is what I don’t need.
Could you give me more advice please? Thanks!