Disable dropout in a model.
(model: torch.nn.Module)
| 97 | |
| 98 | |
| 99 | def disable_dropout(model: torch.nn.Module): |
| 100 | """Disable dropout in a model.""" |
| 101 | for module in model.modules(): |
| 102 | if isinstance(module, torch.nn.Dropout): |
| 103 | module.p = 0 |
| 104 | |
| 105 | |
| 106 | def print_gpu_memory(rank: int = None, message: str = ''): |