def METHOD_NAME(self, src, tgt, seg, soft_tgt=None):
    """
    Args:
        src: [batch_size x choices_num x seq_length]
        tgt: [batch_size]
        seg: [batch_size x choices_num x seq_length]
    """
    choices_num = src.shape[1]
    src = src.view(-1, src.size(-1))
    seg = seg.view(-1, seg.size(-1))
    # Embedding.
    emb = self.embedding(src, seg)
    # Encoder.
    output = self.encoder(emb, seg)
    output = self.dropout(output)
    logits = self.output_layer(output[:, 0, :])
    reshaped_logits = logits.view(-1, choices_num)
    if tgt is not None:
        loss = nn.NLLLoss()(nn.LogSoftmax(dim=-1)(reshaped_logits), tgt.view(-1))
        return loss, reshaped_logits
    else:
        return None, reshaped_logits