我可以让我的自定义pytorch模块在调用train()或eval()时有不同的行为吗？

import torch
import torch.nn as nn

class MyModule(nn.Module):
    def __init__(self):
        super(MyModule, self).__init__()
        # 定义模型的结构

    def forward(self, x):
        # 在forward方法中定义模型的前向传播逻辑
        if self.training:
            # 在训练阶段执行的操作
            # ...
            return train_output
        else:
            # 在评估阶段执行的操作
            # ...
            return eval_output

# 创建模型实例
model = MyModule()

# 在训练阶段调用train()
model.train()
train_output = model(input_data)

# 在评估阶段调用eval()
model.eval()
eval_output = model(input_data)