def get_attention(attn_mat, idx=0, layer=0, head=0):
    """
    Get one slice of the attention map.

    Parameters
    ----------
    attn_mat : Tensor
        attn_mat is numpy array in the shape of [S, N, H, d, d], where
        S is the total number of data samples,
        N is the layer number in the attention mechanism,
        H is the head number in the attention mechanism, and
        d is the attention dimension in each head.
    idx : int, optional
        Index of the input material. The default is 0.
    layer : int, optional
        Layer number in the attention mechanism. The default is 0.
    head : int, optional
        Head number in the attention mechanism. The default is 0.

    Returns
    -------
    attn : Tensor

    """
    attn_mat = attn_mat
    assert len(attn_mat.shape) == 5, "input attn_map is of the wrong shape"
    attn = attn_mat[idx, layer, head, :, :]
    return attn