code: def get_subset_g(g, mask, num_rels, bidirected=False):
      src, dst = g.edges()
      sub_src = src[mask]
      sub_dst = dst[mask]
      sub_rel = g.edata["etype"][mask]
  
      if bidirected:
          sub_src, sub_dst = torch.cat([sub_src, sub_dst]), torch.cat(
              [sub_dst, sub_src]
          )
          sub_rel = torch.cat([sub_rel, sub_rel + num_rels])
  
      sub_g = dgl.graph((sub_src, sub_dst), num_nodes=g.num_nodes())
      sub_g.edata[dgl.ETYPE] = sub_rel
      return sub_g
  
imports: import dgl
  import numpy
  import torch
  import torch.nn
  import torch.nn.functional
  import tqdm
  from dgl.data.knowledge_graph import FB15k237Dataset
  from dgl.dataloading import GraphDataLoader
  from dgl.nn.pytorch import RelGraphConv
label_desc: Use image node_id 1 for calling a global function with example usage: get_subset_g(g, mask, num_rels, bidirected) and returns: sub_g