code: def step(
      self,
      step: int,
      lprobs,
      scores: Optional[Tensor],
      prev_output_tokens: Optional[Tensor] = None,
      original_batch_idxs: Optional[Tensor] = None,
  ):
      n_channels = len(lprobs)
      bsz, beam_size, _ = lprobs[0].size()
  
      lprobs_list = []
      if step == 0:
          # at the first step all hypotheses are equally likely, so use
          # only the first beam
          for i in range(n_channels):
              lprobs_list.append(
                  lprobs[i][:, ::beam_size, :].contiguous()
              )
      else:
          # make probs contain cumulative scores for each hypothesis
          assert scores is not None
          for i in range(n_channels):
              lprobs_list.append(
                  lprobs[i] + scores[:, :, step - 1, i].unsqueeze(-1)
              )
  
      topk_sum_values, topk_sum_indices, topk_sum_beams = topk_sum(
          lprobs_list, k=beam_size * 2
      )
  
      beams_buf = topk_sum_beams
      indices_buf = topk_sum_indices
      scores_buf = (
          torch.tensor(
              [
                  lprobs_list[i][i_batch, i_beam, i_index]
                  for i in range(len(lprobs_list))
                  for i_batch in range(bsz)
                  for i_beam, i_index in zip(
                      beams_buf[i_batch], indices_buf[i_batch, :, i]
                  )
              ]
          )
          .view(len(lprobs_list), bsz, -1)
          .permute(1, 2, 0)
          .to(lprobs_list[0].device)
      )
  
      # At this point, beams_buf and indices_buf are single-dim and contain relative indices
      return scores_buf, indices_buf, beams_buf
  
imports: from typing import Dict, Optional
  import torch
  import torch.nn
  from torch import Tensor
label_desc: Use image node_id 2 for calling the ContiguousMultichannelBeamSearch obj's underlying member method code with example usage: obj.step(step, lprobs, scores, prev_output_tokens, original_batch_idxs) and returns: scores_buf,  indices_buf,  beams_buf