Running training_step of submodules

[pchalasani]

I have a LightningModule M which contains module A and module B as sub-modules.
A,B are instantiations of the same model-type (via the same params), and this model-type has a complex training_step fn.
In the main module M I have a training_step that does something like:

loss = mdl_A.training_step(batches_for_A...) + mdl_B.training_step(batches_for_B...)

And I do something similar with validation_step.

This is a highly simplified description, but in my setting I am seeing signs that the two training_step invocations on the submodules are not independent, and the mdl_B training_step(...) is somehow influenced by the earlier call to mdl_A.training_step(...).

Is there some best practice I should be aware of when trying to do something like that above?