def train_one_epoch(epoch, dataloader, model, criterion, optimizer, device, writer):
    print('train epoch {}'.format(epoch))
    model.train()
    for (idx, (pre_img, post_img, targets)) in enumerate(dataloader):
        (pre_img, post_img, targets) = (pre_img.to(device), post_img.to(device), targets.to(device))
        outputs = model(pre_img, post_img)
        loss = criterion(outputs, targets)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        print('train-epoch:{} [{}/{}], loss: {:5.3}'.format(epoch, (idx + 1), len(dataloader), loss.item()))
        writer.add_scalar('train/loss', loss.item(), ((len(dataloader) * epoch) + idx))
def train_one_epoch(epoch, dataloader, model, criterion, optimizer, device, writer):
    print('train epoch {}'.format(epoch))
    model.train()
    for (idx, (pre_img, post_img, targets)) in enumerate(dataloader):
        (pre_img, post_img, targets) = (pre_img.to(device), post_img.to(device), targets.to(device))
        outputs = model(pre_img, post_img)
        loss = criterion(outputs, targets)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        print('train-epoch:{} [{}/{}], loss: {:5.3}'.format(epoch, (idx + 1), len(dataloader), loss.item()))
        writer.add_scalar('train/loss', loss.item(), ((len(dataloader) * epoch) + idx))