code: def top_k_logits(logits, k):
      
      if k == 0:
          return logits
      else:
          values = torch.topk(logits, k)[0]
          batch_mins = values[:, -1].view(-1, 1).expand_as(logits)
          return torch.where(logits < batch_mins, torch.ones_like(logits) * -1e10, logits)
docs: Masks everything but the k top entries as -infinity (1e10).
      Used to mask logits such that e^-infinity -> 0 won't contribute to the
      sum of the denominator.