import torch def mutate_remove_dropout(model): for module in model.modules(): if isinstance(module, torch.nn.Dropout): module._backup_p = module.p module.p = 0