title: Batch-wise beam search in pytorch
body: <p>I'm trying to implement a beam search decoding strategy in a text generation model. This is the function that I am using to decode the output probabilities.</p>
  <pre><code>def beam_search_decoder(data, k):
      sequences = [[list(), 0.0]]
      # walk over each step in sequence
      for row in data:
          all_candidates = list()
          for i in range(len(sequences)):
              seq, score = sequences[i]
              for j in range(len(row)):
                  candidate = [seq + [j], score - torch.log(row[j])]
                  all_candidates.append(candidate)
          # sort candidates by score
          ordered = sorted(all_candidates, key=lambda tup:tup[1])
          sequences = ordered[:k]
      return sequences
  </code></pre>
  <p>Now you can see this function is implemented with batch_size 1 in mind. Adding another loop for batch size would make the algorithm <code>O(n^4)</code>. It is slow as it is now. Is there any way to improve the speed of this function. My model output is usually of the size <code>(32, 150, 9907)</code> which follows the format <code>(batch_size, max_len, vocab_size)</code></p>