code: def forward(self, words, mapper):
      out = self.dropout(self.base_model(words)[0])
      out = torch.matmul(mapper.float().cuda().transpose(1, 2), out)
      final1 = torch.matmul(out, self.linear.weight)
      final2 = torch.einsum(
          "bnh,hg,bmg->bnm", out, self.bilinear.weight, final1
      )
      root_score = torch.matmul(out, self.root)
      final2 = final2[:, 1:-1, 1:-1]
      N = final2.shape[1]
      final2[:, torch.arange(N), torch.arange(N)] += root_score[:, 1:-1]
      return final2
  
file_path: pytorch-struct/examples/bert-dep.py
imports: import torchtext
  import torch
  from torch import nn
  from torch_struct import DependencyCRF
  from torch_struct.data import SubTokenizedField, ConllXDataset, TokenBucket
  from torchtext.data import RawField, BucketIterator
  from pytorch_transformers import BertModel, BertTokenizer, AdamW, WarmupLinearSchedule
label_desc: Use image node_id 2 for calling the Model obj's underlying member method code with example usage: obj.forward(words, mapper) and returns: final2