model.train() and model.eval() sets the calling nn.Module's and its children’s modules training attribute to True and False respectively.
If you need to invoke functions based on training or testing mode, you can simply use the network’s training attribute.
def forward(self, x):
    if self.training:
        out = self.train_module(x)
    else:
        out = self.eval_module(x)
    return out