def symmetrize_matrix(A):
    """
    Symmetrizes a matrix :
    If shape is (a,b,c) will symmetrize by considering a is batch size
    """
    Af = A.triu(0) + A.triu(1).transpose(-2,-1)
    return Af