How to optimize two models in a class?

[KasenYoung]

Thanks a lot for your great work. Now I have a class Model which has two little models.

class Model(nn.Module):
       def __init__(self):
             super().__init__()
             self.model1 = xxx
             self.model2 = xxx

When using accelerate package and deepspeed, it only supports one optimizer. So I have to place the parameters of the two models into one optimizer:

optimizer = optimizer(
      params_to_optimize=[model1.parameters(), model2.parameters()],
      ...,
)

My question is : when I optimize just one model, what should I do? For example, I want to optimize model1, what should I do to prevent model2 to be optimized? Thank you very much for your help.

[hwchen2017]

Hi @KasenYoung , you can only pass the parameters of model1 to the optimizer like

optimizer = optimizer(
      params_to_optimize=[Model.model1.parameters()],
      ...,
)