| 86 | """ |
| 87 | |
| 88 | def __init__(self, args, model, optimizer, grad_accum_count=1, n_gpu=1, gpu_rank=1, report_manager=None): |
| 89 | # Basic attributes. |
| 90 | self.args = args |
| 91 | self.save_checkpoint_steps = args.save_checkpoint_steps |
| 92 | self.model = model |
| 93 | self.optimizer = optimizer |
| 94 | self.grad_accum_count = grad_accum_count |
| 95 | self.n_gpu = n_gpu |
| 96 | self.gpu_rank = gpu_rank |
| 97 | self.report_manager = report_manager |
| 98 | |
| 99 | self.loss = torch.nn.BCELoss(reduction='none') |
| 100 | assert grad_accum_count > 0 |
| 101 | # Set models in training mode. |
| 102 | if model: |
| 103 | self.model.train() |
| 104 | |
| 105 | def train(self, train_iter_fct, train_steps, valid_iter_fct=None, valid_steps=-1): |
| 106 | """ |