def _train_batch(loader, model, optimizer, loss_fun):
    """Train a batch and return the loss"""
    running_loss = 0
    total = 0
    for batch_idx, (target, feat) in enumerate(loader):
        pred = model(feat)
        optimizer.zero_grad()
        loss = loss_fun(pred.flatten(), target)
        loss.backward()
        optimizer.step()
        running_loss += loss.item()
        total = batch_idx

    return running_loss / (total + 1)