Skip to content

Instantly share code, notes, and snippets.

@penglin03

penglin03/test.py

Created May 1, 2025 08:25
Show Gist options

  • Select an option

    Learn more about clone URLs
  • Save penglin03/c2c1db618cd5e90fe5821ede30d7533b to your computer and use it in GitHub Desktop.

Select an option

Learn more about clone URLs
Save penglin03/c2c1db618cd5e90fe5821ede30d7533b to your computer and use it in GitHub Desktop.
Download ZIP
Raw
test.py
import dgl
from dgl import DGLGraph
import torch
from dgl import DropEdge
def drop_edges_in_blocks(blocks, drop_prob=0.2):
new_blocks = []
for block in blocks:
kept_edges = {}
for etype in block.etypes:
num_edges_of_etype = block.num_edges(etype)
mask = torch.rand(num_edges_of_etype) > drop_prob
kept_edges[etype] = torch.arange(num_edges_of_etype)[mask]
block_with_drop = dgl.edge_subgraph(block, kept_edges, relabel_nodes=False)
new_blocks.append(block_with_drop)
return new_blocks
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment