(save_every: int, total_epochs: int, batch_size: int, snapshot_path: str = "snapshot.pt")
| 93 | |
| 94 | |
| 95 | def main(save_every: int, total_epochs: int, batch_size: int, snapshot_path: str = "snapshot.pt"): |
| 96 | ddp_setup() |
| 97 | dataset, model, optimizer = load_train_objs() |
| 98 | train_data = prepare_dataloader(dataset, batch_size) |
| 99 | trainer = Trainer(model, train_data, optimizer, save_every, snapshot_path) |
| 100 | trainer.train(total_epochs) |
| 101 | destroy_process_group() |
| 102 | |
| 103 | |
| 104 | if __name__ == "__main__": |
no test coverage detected