def METHOD_NAME(self, src, tgt, seg, soft_tgt=None):
    """
    Args:
        src: [batch_size x seq_length]
        tgt: [batch_size]
        seg: [batch_size x seq_length]
    """
    # Embedding.
    emb = self.embedding(src, seg)
    # Encoder.
    output = self.encoder(emb, seg)
    # Target.
    output = pooling(output, seg, self.pooling_type)
    output = torch.tanh(self.output_layer_1(output))
    logits = self.output_layer_2(output)
    if tgt is not None:
        loss = nn.MSELoss()(logits.view(-1), tgt.view(-1))
        return loss, logits
    else:
        return None, logits