code: def train(train_iter, val_iter, model):
      opt = AdamW(model.parameters(), lr=1e-4, eps=1e-8)
      scheduler = WarmupLinearSchedule(
          opt, warmup_steps=20, t_total=2500
      )
      model.train()
      losses = []
      for i, ex in enumerate(train_iter):
          opt.zero_grad()
          words, mapper, _ = ex.word
          label, lengths = ex.head
          batch, _ = label.shape
  
          # Model
          final = model(words.cuda(), mapper)
          for b in range(batch):
              final[b, lengths[b] - 1 :, :] = 0
              final[b, :, lengths[b] - 1 :] = 0
  
          if not lengths.max() <= final.shape[1] + 1:
              print("fail")
              continue
          dist = DependencyCRF(final, lengths=lengths)
  
          labels = dist.struct.to_parts(label, lengths=lengths).type_as(
              final
          )
          log_prob = dist.log_prob(labels)
  
          loss = log_prob.sum()
          (-loss).backward()
          torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
  
          opt.step()
          scheduler.step()
          losses.append(loss.detach())
          if i % 50 == 1:
              print(-torch.tensor(losses).mean(), words.shape)
              losses = []
          if i % 600 == 500:
              validate(val_iter)
  
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 4 for calling a global function with example usage: train(train_iter, val_iter, model) without return types