def _train(self, loader):
    self.model.train()
    (bar, metric) = (progress_bar(loader), AttachmentMetric())
    for (i, batch) in enumerate(bar, 1):
        (words, texts, *feats, arcs, rels) = batch
        word_mask = words.ne(self.args.pad_index)
        mask = (word_mask if (len(words.shape) < 3) else word_mask.any((- 1)))
        mask[:, 0] = 0
        (s_arc, s_rel) = self.model(words, feats)
        loss = self.model.loss(s_arc, s_rel, arcs, rels, mask, self.args.partial)
        loss = (loss / self.args.update_steps)
        loss.backward()
        nn.utils.clip_grad_norm_(self.model.parameters(), self.args.clip)
        if ((i % self.args.update_steps) == 0):
            self.optimizer.step()
            self.scheduler.step()
            self.optimizer.zero_grad()
        (arc_preds, rel_preds) = self.model.decode(s_arc, s_rel, mask)
        if self.args.partial:
            mask &= arcs.ge(0)
        if (not self.args.punct):
            mask.masked_scatter_(mask, (~ mask.new_tensor([ispunct(w) for s in texts for w in s])))
        metric(arc_preds, rel_preds, arcs, rels, mask)
        bar.set_postfix_str(f'lr: {self.scheduler.get_last_lr()[0]:.4e} - loss: {loss:.4f} - {metric}')
    logger.info(f'{bar.postfix}')
def _train(self, loader):
    self.model.train()
    (bar, metric) = (progress_bar(loader), AttachmentMetric())
    for (i, batch) in enumerate(bar, 1):
        (words, texts, *feats, arcs, rels) = batch
        word_mask = words.ne(self.args.pad_index)
        mask = (word_mask if (len(words.shape) < 3) else word_mask.any((- 1)))
        mask[:, 0] = 0
        (s_arc, s_rel) = self.model(words, feats)
        loss = self.model.loss(s_arc, s_rel, arcs, rels, mask, self.args.partial)
        loss = (loss / self.args.update_steps)
        loss.backward()
        nn.utils.clip_grad_norm_(self.model.parameters(), self.args.clip)
        if ((i % self.args.update_steps) == 0):
            self.optimizer.step()
            self.scheduler.step()
            self.optimizer.zero_grad()
        (arc_preds, rel_preds) = self.model.decode(s_arc, s_rel, mask)
        if self.args.partial:
            mask &= arcs.ge(0)
        if (not self.args.punct):
            mask.masked_scatter_(mask, (~ mask.new_tensor([ispunct(w) for s in texts for w in s])))
        metric(arc_preds, rel_preds, arcs, rels, mask)
        bar.set_postfix_str(f'lr: {self.scheduler.get_last_lr()[0]:.4e} - loss: {loss:.4f} - {metric}')
    logger.info(f'{bar.postfix}')