A context manager that places a model into training mode and restores the previous mode on exit.
(model, mode=True)
| 70 | |
| 71 | @contextmanager |
| 72 | def train_mode(model, mode=True): |
| 73 | """A context manager that places a model into training mode and restores |
| 74 | the previous mode on exit.""" |
| 75 | modes = [module.training for module in model.modules()] |
| 76 | try: |
| 77 | yield model.train(mode) |
| 78 | finally: |
| 79 | for i, module in enumerate(model.modules()): |
| 80 | module.training = modes[i] |
| 81 | |
| 82 | |
| 83 | def eval_mode(model): |