Dropout and Batch-Norm already have this kind of behavior. In other words what a Dropout or Batch-Norm module outputs when it is in eval mode is different from when it is in train mode. How can I achieve something similar with my custom module?
You could use the internal
Here is a dummy example:
class MyModel(nn.Module): def __init__(self): super(MyModel, self).__init__() self.fc = nn.Linear(10, 10) def forward(self, x): x = self.fc(x) if self.training: x = x * 1000 return x model = MyModel() x = torch.randn(1, 10) output = model(x) print(output) > tensor([[ -151.6117, 20.6451, -589.1161, -120.6478, 395.1652, -950.3046, -1062.1073, 973.9295, 61.4954, -412.5521]], grad_fn=<MulBackward0>) model.eval() output = model(x) print(output) > tensor([[-0.1516, 0.0206, -0.5891, -0.1206, 0.3952, -0.9503, -1.0621, 0.9739, 0.0615, -0.4126]], grad_fn=<AddmmBackward>)
Thank you. That seems very nice and simple.