import torch


def uv_flow_loss(uv_flow_pred, uv_flow_gt_w_mask):
    """
    Input:
      uv_flow_pred: B x 2 x H x W
      uv_flow_gt_w_mask: B x 3 x H x W
    """
    # We only have info for the uv_flow for the points where mask is one
    mask = uv_flow_gt_w_mask[:,2,:,:].unsqueeze(1)
    uv_flow_gt = uv_flow_gt_w_mask[:,:2,:,:]
    return torch.nn.L1Loss()(uv_flow_pred*mask, uv_flow_gt*mask)