I would advise you to use a separate optimizer for each model/submodel. It gives you more control and you don’t have to syncronize optimizer.zero_grad() and optimizer.step() calls (otherwise, you have to train all the models at the same time). If you decide to use a single optimizer, you can define a different learning policy for each model through parameter groups that are explained in documentation.