def remove_self_loops(edge_index, edge_sqdist, edge_weight=None):
    """
    Removes all self-loop in the graph.

    Args:
        edge_index (`torch.LongTensor`): indices of graph's edges.
        edge_attr (`torch.FloatTensor`): attributes of graph's edges.

    Returns:
        (`torch.LongTensor`): indices of graph's edges.
        (`torch.FloatTensor`): attributes of graph's edges.
    """
    mask = edge_index[0] != edge_index[1]

    if edge_weight is None:
        return edge_index[..., mask], edge_sqdist[mask]

    return edge_index[..., mask], edge_sqdist[mask], edge_weight[mask]