def calc_norm_MSE_loss(Y_pred, Y, u):
    """
    Calculates the normalized MSE loss using ground truth Y, visit density
    predictions Y_pred, and loss weights u.

    Formula for normalized MSE loss:
        Fro-norm( diag(u)(Y - Y_pred) )^2 / Fro-norm( diag(u)(Y - mean(Y)) )^2

    The formula aligns with the Identification Problem's metric used in the
    study. mean(Y) is row-wise mean of Y.

    Args:
        Y_pred -- (torch.Tensor) T by J tensor of visit density predictions.
        Y -- (torch.Tensor) T by J tensor of ground truth future density
            predictions.
        u -- (torch.Tensor) T-length vector of weights for calculating loss.
    """
    # Normalizer is diag(u)(Y - mean(Y)), obtained by making 1d data['u']
    # 2d and multiplying element-wise
    normalizer = u.unsqueeze(-1) * (Y - Y.mean(1, keepdim=True))

    # Loss due to predictions X = Y
    loss_preds = u.unsqueeze(-1) * (Y - Y_pred)

    # Calculate Frobenius norm of both losses and divide to get normalized
    # MSE loss
    return (loss_preds.norm(2) ** 2) / (normalizer.norm(2) ** 2)