(self, batch)
| 432 | return x |
| 433 | |
| 434 | def shared_step(self, batch): |
| 435 | x = self.get_input(batch, self.first_stage_key) |
| 436 | loss, loss_dict = self(x) |
| 437 | return loss, loss_dict |
| 438 | |
| 439 | def training_step(self, batch, batch_idx): |
| 440 | for k in self.ucg_training: |
no test coverage detected