Params:
                input_ids :: [bsz, len]
                target :: [bsz, len]
            Returns:
                tuple(softmax_output, new_mems) where:
                    new_mems: list (num layers) of hidden states at the entry of each layer
                        shape :: [mem_len, bsz, self.config.d_model] :: Warning: shapes are transposed here w. regards to input_ids
                    softmax_output: output of the (adaptive) softmax:
                        if target is None:
                            Negative log likelihood of shape :: [bsz, len] 
                        else:
                            log probabilities of tokens, shape :: [bsz, len, n_tokens]

def forward(self, input_ids, target=None, mems=None):
        """ Params:
                input_ids :: [bsz, len]
                target :: [bsz, len]
            Returns:
                tuple(softmax_output, new_mems) where:
                    new_mems: list (num layers) of hidden states at the entry of each layer
                        shape :: [mem_len, bsz, self.config.d_model] :: Warning: shapes are transposed here w. regards to input_ids
                    softmax_output: output of the (adaptive) softmax:
                        if target is None:
                            Negative log likelihood of shape :: [bsz, len] 
                        else:
                            log probabilities of tokens, shape :: [bsz, len, n_tokens]
        """
        bsz = input_ids.size(0)
        tgt_len = input_ids.size(1)

        last_hidden, new_mems = self.transformer(input_ids, mems)

        pred_hid = last_hidden[:, -tgt_len:]
        if self.sample_softmax > 0 and self.training:
            assert self.config.tie_weight
            logit = sample_logits(self.transformer.word_emb, self.out_layer.bias, target, pred_hid, self.sampler)
            softmax_output = -F.log_softmax(logit, -1)[:, :, 0]
        else:
            softmax_output = self.crit(pred_hid.view(-1, pred_hid.size(-1)), target)
            if target is None:
                softmax_output = softmax_output.view(bsz, tgt_len, -1)
            else:
                softmax_output = softmax_output.view(bsz, tgt_len)

        # We transpose back
        return (softmax_output, new_mems)