code: def forward(self, input_ids, target=None, mems=None):
          
          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)
docs: 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]