def METHOD_NAME(x_hat, x, mask, name=False):
    x = x[:, None] if len(x.shape) == 2 else x  # add channel dim
    mask = mask[:, None] if len(mask.shape) == 2 else mask  # add channel dim
    assert len(x.shape) == len(mask.shape)
    x = x * mask
    x_hat = x_hat * mask
    if name == 'vpred':
        loss = F.binary_cross_entropy_with_logits(x_hat, x, reduction='sum')
    else:
        loss = F.mse_loss(x_hat, x, reduction='sum')
    loss = loss / mask.sum()
    loss_dict = {"loss_{}".format(name): loss}
    return loss_dict