I think this is working! However, I had a doubt. In my implementation, I am doing the following:
Create a ModelMerger class where I pass a list of models, and I initialize an optimizer with what @ptrblck mentioned above.
Then I iteratively do:

Forward propagation of ModelA, computation of loss function, backward propagation, accumulation of gradients following idea in https://discuss.pytorch.org/t/how-to-get-and-or-set-gradient-data-during-training 15.
Extraction of layer A2 to use it as input for ModelB.
Forward propagation of ModelB using A2 as inputs, and computation of loss function, accumulation of gradients.

Once each model finishes the steps above, I do a self.optimizer.step(). I just do a
optimizer = optim.SGD(list(modelA.parameters()) + list(modelB.parameters()), lr=1e-3)

Is this correct? It seems to me that pytorch is so smart that figures out everything. I wonder if I should sum the losses before doing loss.backward() or running them independently is fine? Would that change anything?
@ptrblck, @alex.veuthey I would appreciate any input from you guys :), and thanks for your help!