code: def drop_edge_weighted(
      edge_index, edge_weights, p: float, threshold: float
  ):
      """Returns the new edge index after probabilistically dropping edges from `edge_index`
      Args:
          edge_index (torch.Tensor): Tensor of shape [2, n_edges]
          edge_weights (torch.Tensor): Tensor of shape [n_edges]
          p (float): Probability multiplier
          threshold (float): Upper bound probability of dropping an edge
      """
      edge_weights = edge_weights / edge_weights.mean() * p
      edge_weights = edge_weights.where(
          edge_weights < threshold,
          torch.ones_like(edge_weights) * threshold,
      )
      sel_mask = torch.bernoulli(1.0 - edge_weights).to(torch.bool)
  
      return edge_index[:, sel_mask]
  
file_path: DIG/dig/sslgraph/utils/adaptive.py
imports: import torch
  from torch.utils.data import random_split
  from torch_geometric.utils import degree, remove_self_loops, get_laplacian, to_undirected, to_dense_adj, to_networkx, degree, to_undirected
  from torch_scatter import scatter
  import networkx
label_desc: Use image node_id 6 for calling a global function with example usage: drop_edge_weighted(edge_index, edge_weights, p, threshold) and returns: edge_index