(rank, args, model, device, dataset, dataloader_kwargs)
| 5 | |
| 6 | |
| 7 | def train(rank, args, model, device, dataset, dataloader_kwargs): |
| 8 | torch.manual_seed(args.seed + rank) |
| 9 | |
| 10 | train_loader = torch.utils.data.DataLoader(dataset, **dataloader_kwargs) |
| 11 | |
| 12 | optimizer = optim.SGD(model.parameters(), lr=args.lr, momentum=args.momentum) |
| 13 | for epoch in range(1, args.epochs + 1): |
| 14 | train_epoch(epoch, args, model, device, train_loader, optimizer) |
| 15 | |
| 16 | |
| 17 | def test(args, model, device, dataset, dataloader_kwargs): |
nothing calls this directly
no test coverage detected