import torch


def batch_matrix_norm(matrix, norm_order=2):
    """ normalization of the matrix
    Args:
        matrix: torch.Tensor. Expected shape [batch, *]
        norm_order: int. Order of normalization. 

    Returns:
        normed matrix: torch.Tensor. 
    """
    return torch.norm(matrix, p=norm_order, dim=[1,2])