| 45 | |
| 46 | |
| 47 | def train(epoch): |
| 48 | epoch_loss = 0 |
| 49 | for iteration, batch in enumerate(training_data_loader, 1): |
| 50 | input, target = batch[0].to(device), batch[1].to(device) |
| 51 | |
| 52 | optimizer.zero_grad() |
| 53 | loss = criterion(model(input), target) |
| 54 | epoch_loss += loss.item() |
| 55 | loss.backward() |
| 56 | optimizer.step() |
| 57 | |
| 58 | print("===> Epoch[{}]({}/{}): Loss: {:.4f}".format(epoch, iteration, len(training_data_loader), loss.item())) |
| 59 | |
| 60 | print("===> Epoch {} Complete: Avg. Loss: {:.4f}".format(epoch, epoch_loss / len(training_data_loader))) |
| 61 | |
| 62 | |
| 63 | def test(): |