(self, device: torch.device)
| 105 | return optim.Adam(self.parameters(), lr=self.cfg.lr) |
| 106 | |
| 107 | def export(self, device: torch.device) -> ModelExports: |
| 108 | return self.model.export( |
| 109 | self.batch.to(device), |
| 110 | self.flows.to(device), |
| 111 | self.global_step, |
| 112 | ) |