(self, source, targets)
| 38 | self.model = DDP(model, device_ids=[gpu_id]) |
| 39 | |
| 40 | def _run_batch(self, source, targets): |
| 41 | self.optimizer.zero_grad() |
| 42 | output = self.model(source) |
| 43 | loss = F.cross_entropy(output, targets) |
| 44 | loss.backward() |
| 45 | self.optimizer.step() |
| 46 | |
| 47 | def _run_epoch(self, epoch): |
| 48 | b_sz = len(next(iter(self.train_data))[0]) |