code: def prepare_edges_for_metis(
      edge_index: torch.Tensor,
  ) -> torch.LongTensor:
      """Prepare the edge index for METIS partitioning to prevent segfaults."""
      # remove self-loops
      mask = edge_index[0] != edge_index[1]
      edge_index = edge_index[:, mask]
  
      # add inverse edges and remove duplicates
      return torch.cat([edge_index, edge_index.flip(0)], dim=-1).unique(
          dim=1
      )
  
file_path: pykeen/src/pykeen/nn/node_piece/utils.py
imports: import logging
  from typing import Collection, Mapping, Optional
  import numpy
  import torch
  from tqdm.auto import tqdm
label_desc: Use image node_id 3 for calling a global function with example usage: prepare_edges_for_metis(edge_index) and returns: torch