Solved by ptrblck in post #2 
              
              
                You could use the internal self.training attribute. 
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:
           …