def forward(self, src, tgt, seg):
        """
        Args:
            src: [batch_size x seq_length]
            tgt: [batch_size x seq_length]
            seg: [batch_size x seq_length]
        Returns:
            loss: Sequence labeling loss.
            logits: Output logits.
        """
        # Embedding.
        emb = self.embedding(src, seg)
        # Encoder.
        output = self.encoder(emb, seg)

        # Target.
        logits = self.output_layer(output)
        if self.crf_target:
            tgt_mask = seg.type(torch.uint8)
            pred = self.crf.decode(logits, mask=tgt_mask)
            for j in range(len(pred)):
                while len(pred[j]) < self.seq_length:
                    pred[j].append(self.labels_num - 1)
            pred = torch.tensor(pred).contiguous().view(-1)
            if tgt is not None:
                loss = -self.crf(F.log_softmax(logits, 2), tgt, mask=tgt_mask, reduction='mean')
                return loss, pred
            else:
                return None, pred
        else:
            tgt_mask = seg.contiguous().view(-1).float()
            logits = logits.contiguous().view(-1, self.labels_num)
            pred = logits.argmax(dim=-1)
            if tgt is not None:
                tgt = tgt.contiguous().view(-1, 1)
                one_hot = torch.zeros(tgt.size(0), self.labels_num). \
                    to(torch.device(tgt.device)). \
                    scatter_(1, tgt, 1.0)
                numerator = -torch.sum(nn.LogSoftmax(dim=-1)(logits) * one_hot, 1)
                numerator = torch.sum(tgt_mask * numerator)
                denominator = torch.sum(tgt_mask) + 1e-6
                loss = numerator / denominator
                return loss, pred
            else:
                return None, pred

        Args:
            src: [batch_size x seq_length]
            tgt: [batch_size x seq_length]
            seg: [batch_size x seq_length]
        Returns:
            loss: Sequence labeling loss.
            logits: Output logits.
        
finetune/run_ner.py