The way to do this is the following (working only on nightly pytorch build right now and will be in 1.8 when it comes out):
a_loss.backward(inputs=tm.a.parameters(), retain_graph=True)
b_loss.backward(inputs=tm.b.parameters())

tm.a_optimizer.step()  
tm.b_optimizer.step()

Hope this helps!
Sorry for the confusion earlier!