(self, batch, batch_idx)
| 557 | return out |
| 558 | |
| 559 | def training_step(self, batch, batch_idx): |
| 560 | if self.controller is not None: |
| 561 | self.controller.between_steps() |
| 562 | if isinstance(batch, list): |
| 563 | train_batch = batch[0] |
| 564 | train2_batch = batch[1] |
| 565 | loss_train, loss_dict = self.shared_step(train_batch) |
| 566 | loss_train2, _ = self.shared_step(train2_batch) |
| 567 | loss = loss_train + loss_train2 |
| 568 | else: |
| 569 | train_batch = batch |
| 570 | loss, loss_dict = self.shared_step(train_batch) |
| 571 | |
| 572 | self.log_dict(loss_dict, prog_bar=True, |
| 573 | logger=True, on_step=True, on_epoch=True) |
| 574 | |
| 575 | self.log("global_step", self.global_step, |
| 576 | prog_bar=True, logger=True, on_step=True, on_epoch=False) |
| 577 | |
| 578 | if self.use_scheduler: |
| 579 | lr = self.optimizers().param_groups[0]['lr'] |
| 580 | self.log('lr_abs', lr, prog_bar=True, logger=True, on_step=True, on_epoch=False) |
| 581 | return loss |
| 582 | |
| 583 | def shared_step(self, batch, **kwargs): |
| 584 | x, c, mask = self.get_input_withmask(batch, **kwargs) |
nothing calls this directly
no test coverage detected