def train_model(model, optim, train_dl, loss_fn, use_gpu=False):
    """
    Controls training loop for MatchNN
    :param model: pytorch model instance
    :param optim: pytorch optimizer
    :param train_dl: training dataloader on top of MatchDataset
    :param loss_fn: loss function to train model
    :param use_gpu: boolean for whether GPU will be used
    :return: loss
    """
    model.train()
    total = 0
    sum_loss = 0
    for x1, x2, y in train_dl:
        if use_gpu:
            x1, x2, y = x1.cuda(), x2.cuda(), y.cuda()
        batch = y.shape[0]
        output = model(x1, x2)
        loss = loss_fn(output, y)
        optim.zero_grad()
        loss.backward()
        optim.step()
        total += batch
        sum_loss += batch * (loss.item())
    return sum_loss / total