Manually call model.eval() and model.train() inside the training loop

[rwbfd]

How to manually call model.eval() or model.train() inside the lightning module? I happen to have several models and not all of them need to be updated during each forward pass. Thanks!

[kkirtac]

You can use self.freeze() and self.unfreeze().

[rwbfd]

Hi thanks for the reply. However, what if I have two models and I need to freeze one while unfreeze the other one?

[kkirtac]

That shouldn't make a difference. Here in this example there are 2 models, i.e., encoder and decoder. If your encoder and decoder models are type of nn.Module lets say, you can call anyone like self.encoder.train() or self.encoder.eval() and the same applies for the self.decoder as well. If they are type of pl.LightningModule then you can use freeze() and unfreeze() as examplified above.