import torch


def compute_fw_loss(flo_diff, mask=None):
    """
        compute forward-backward flow consistency loss
    :param flo_diff: difference of flow and warped-flow
    :param mask: disocclusion mask
    :return: loss
    """
    if mask == None:
        return torch.mean(torch.abs(flo_diff))
    else:
        return torch.abs(flo_diff * mask).mean() / mask.mean()