I am not sure if this is what you want. Maybe you can revert them before creating blocks.
class HomoNeighborSampler(NeighborSampler):
def sample_blocks(self, g, seed_nodes, exclude_eids=None):
output_nodes = seed_nodes
blocks = []
# convert to homogeneous graph first
# NOTE pass the list of feature names in data fields
homo_g = dgl.to_homogeneous(g, ndata=['nfeats'], edata=['efeats'])
for fanout in reversed(self.fanouts):
# sample on the homogeneous graph
frontier = homo_g.sample_neighbors(
seed_nodes, fanout, edge_dir=self.edge_dir, prob=self.prob,
replace=self.replace, output_device=self.output_device,
exclude_edges=exclude_eids)
# revert the subgraph to heterogeneous graph before creating the block
frontier = dgl.to_heterogeneous(frontier, g.ntypes, g.etypes)
eid = frontier.edata[EID]
block = to_block(frontier, seed_nodes)
block.edata[EID] = eid
seed_nodes = block.srcdata[NID]
blocks.insert(0, block)
return seed_nodes, output_nodes, blocks