How the model.train() and model.eval() switch work in Pytorch Lightning?

[lndip]

I am implementing a knowledge distillation (KD) framework. The teacher model architecture has the BatchNorm and Dropout which must be in evaluation mode during the training of the student in the KD process. Is freezing the teacher in the initialization of the LightningModule as followed is enough? Or do I have to explicitly do that use the hook on_train_start()

class BaseDistilledModule(LightningModule): 
    def __init__(self,
                 teacher: Module,
                 student: Module,
                 teacher_ckpt_path: Path)
    self.teacher = teacher
    self.student = student

    self.teacher.load_state_dict(torch.load(teacher_ckpt_path)["state_dict"])
    # Freeze teacher
    self.teacher.eval()
    for param in self.teacher.parameters():
        param.requires_grad = False

[mirsazzathossain]

@lndip Did you get any solution?

I am in a similar kind of scenario. I am trying to implement a loss with learnable weight.

$loss = \alpha *$ some_loss $+ \beta *$ other_loss

Here, $\alpha$ and $\beta$ are learnable params. I am unsure where to put loss.train().