(rank: int, world_size: int, save_every: int, total_epochs: int, batch_size: int)
| 84 | |
| 85 | |
| 86 | def main(rank: int, world_size: int, save_every: int, total_epochs: int, batch_size: int): |
| 87 | ddp_setup(rank, world_size) |
| 88 | dataset, model, optimizer = load_train_objs() |
| 89 | train_data = prepare_dataloader(dataset, batch_size) |
| 90 | trainer = Trainer(model, train_data, optimizer, rank, save_every) |
| 91 | trainer.train(total_epochs) |
| 92 | destroy_process_group() |
| 93 | |
| 94 | |
| 95 | if __name__ == "__main__": |
nothing calls this directly
no test coverage detected