Function main

distributed/ddp-tutorial-series/multigpu_torchrun.py:94–100

(save_every: int, total_epochs: int, batch_size: int, snapshot_path: str = "snapshot.pt")

def main(save_every: int, total_epochs: int, batch_size: int, snapshot_path: str = "snapshot.pt"):
    ddp_setup()
    dataset, model, optimizer = load_train_objs()
    train_data = prepare_dataloader(dataset, batch_size)
    trainer = Trainer(model, train_data, optimizer, save_every, snapshot_path)
    trainer.train(total_epochs)
    destroy_process_group()


if __name__ == "__main__":