import torch


def weighted_norm(input, Re):
    """
    Returns a new tensor whose elements correspond to the squared weighted norm of the input.

    Args:
        input (`torch.FloatTensor`): input tensor with shape (N, D).
        Re (`torch.FloatTensor`): metric tensor with shape (D, D).

    Returns:
        (`torch.FloatTensor`): squared weighted norm with shape (N).
    """
    Re = Re.to(input.device)
    return torch.matmul(torch.matmul(input.transpose(1, 2), Re), input).squeeze()