If you replaced the modules with the functional API calls, you would have to pass the training flag manually:
def forward(self, x):
    x = F.dropout(x, training=self.training)

Could you post the model definition so that we can have a look?