def compute_loss(criterion, outputs, labels, batch_size):
    """
    Helper function to compute the loss. Since this is a pixel-wise
    prediction task we need to reshape the output and ground truth
    tensors into a 2D tensor before passing it in to the loss criterion.
    Args:
      criterion: pytorch loss criterion
      outputs (pytorch tensor): predicted labels from the model
      labels (pytorch tensor): ground truth labels
      batch_size (int): batch size used for training
    Returns:
      pytorch tensor for loss
    """
    loss_out = outputs.transpose(1, 3) \
        .contiguous() \
        .view([batch_size * 128 * 128, 2])
    loss_lab = labels.transpose(1, 3) \
        .contiguous() \
        .view([batch_size * 128 * 128, 2])
    return criterion(loss_out, loss_lab)