def make_nbr_dic(nbr_list):
    """
    Make a dictionary that maps each atom to the indices
    of its neighbors.
    Args:
        nbr_list (torch.LongTensor): nbr list for a geometry
    Returns:
        nbr_dic (dict): dictionary described above
    """

    nbr_dic = {}
    for nbr in nbr_list:
        nbr_0 = nbr[0].item()
        if nbr_0 not in nbr_dic:
            nbr_dic[nbr_0] = []
        nbr_dic[nbr_0].append(nbr[1].item())
    return nbr_dic