code: def get_mask(self, ents, bs, ne, run_device):
      ent_mask = torch.zeros(bs, ne, device=run_device)
      rel_mask = torch.zeros(bs, ne, ne, device=run_device)
      for _b in range(bs):
          ent_mask[_b, : len(ents[_b])] = 1
          rel_mask[_b, : len(ents[_b]), : len(ents[_b])] = 1
      return ent_mask, rel_mask
  
file_path: DeepKE/src/deepke/relation_extraction/document/model.py
imports: import torch
  import torch.nn
  from opt_einsum import contract
  from .losses import ATLoss
  import torch.nn.functional
  from .module import AttentionUNet
  from .prepro import process_long_input
label_desc: Use image node_id 4 for calling the DocREModel obj's underlying member method code with example usage: obj.get_mask(ents, bs, ne, run_device) and returns: ent_mask,  rel_mask