import torch


def bgram(m: torch.Tensor) -> torch.Tensor:
    """
    Return the batched Gram matrix of `m`

    Args:
        m (tensor): a matrix of dim 3, first one is the batch

    Returns:
        The batch of Gram matrix
    """
    assert m.dim() == 4
    m = m.view(m.shape[0], m.shape[1], -1)
    g = torch.einsum('bik,bjk->bij', m, m) / m.shape[2]
    return g