| 34 | |
| 35 | |
| 36 | def train(args, model, device, train_loader, optimizer, epoch): |
| 37 | model.train() |
| 38 | for batch_idx, (data, target) in enumerate(train_loader): |
| 39 | data, target = data.to(device), target.to(device) |
| 40 | optimizer.zero_grad() |
| 41 | output = model(data) |
| 42 | loss = F.nll_loss(output, target) |
| 43 | loss.backward() |
| 44 | optimizer.step() |
| 45 | if batch_idx % args.log_interval == 0: |
| 46 | print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format( |
| 47 | epoch, batch_idx * len(data), len(train_loader.dataset), |
| 48 | 100. * batch_idx / len(train_loader), loss.item())) |
| 49 | if args.dry_run: |
| 50 | break |
| 51 | |
| 52 | |
| 53 | def test(model, device, test_loader): |