def train()

in retrieval_train.py [0:0]

• 39 lines of code
• 7 McCabe index (conditional complexity)

def train(epoch, start_time, model, optimizer, opt_, data_loader):
    """Run through one epoch of model training with the provided data loader."""
    model.train()
    # Initialize meters + timers
    train_loss = 0
    nb_ok = 0
    nb_exs = 0
    nb_losses = 0
    epoch_start = time.time()
    # Run one epoch
    for idx, ex in enumerate(data_loader, 1):
        params = [
            field.cuda(non_blocking=True)
            if opt_.cuda
            else field
            if field is not None
            else None
            for field in ex
        ]
        loss, ok = loss_fn(*model(*params))
        nb_ok += ok
        nb_exs += ex[0].size(0)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        train_loss += loss.sum().item()
        nb_losses += 1
        if idx % opt_.display_iter == 0 or idx == len(data_loader):
            avg_loss = train_loss / nb_losses
            acc = 100 * nb_ok / nb_exs
            elapsed = time.time() - start_time
            logging.info(
                f"train: Epoch = {epoch} | iter = {idx}/{len(data_loader)} | loss = "
                f"{avg_loss:.3f} | batch P@1 = {acc:.2f} % | elapsed time = "
                f"{elapsed:.2f} (s)"
            )
            train_loss = 0
            nb_losses = 0
    epoch_elapsed = time.time() - epoch_start
    logging.info(
        f"train: Epoch {epoch:d} done. Time for epoch = {epoch_elapsed:.2f} (s)"
    )