code: def step(
      self,
      step: int,
      lprobs,
      scores: Optional[Dict[str, Tensor]],
      prev_output_tokens: Optional[Dict[str, Tensor]] = None,
      original_batch_idxs: Optional[Tensor] = None,
  ):
      channels = list(lprobs.keys())
      bsz, beam_size, _ = lprobs[channels[0]].size()
  
      lprobs_list = []
      if step == 0:
          # at the first step all hypotheses are equally likely, so use
          # only the first beam
          for channel in channels:
              lprobs_list.append(
                  lprobs[channel][:, ::beam_size, :].contiguous()
              )
      else:
          # make probs contain cumulative scores for each hypothesis
          assert scores is not None
          for channel in channels:
              lprobs_list.append(
                  lprobs[channel]
                  + scores[channel][:, :, step - 1].unsqueeze(-1)
              )
  
      topk_sum_values, topk_sum_indices, topk_sum_beams = topk_sum(
          lprobs_list, k=beam_size * 2
      )
  
      beams_buf = topk_sum_beams
      scores_buf = {}
      indices_buf = {}
      for i, channel in enumerate(channels):
          indices_buf[channel] = topk_sum_indices[:, :, i]
          scores_buf[channel] = (
              torch.tensor(
                  [
                      lprobs_list[i][i_batch, i_beam, i_index]
                      for i_batch in range(bsz)
                      for i_beam, i_index in zip(
                          beams_buf[i_batch],
                          indices_buf[channel][i_batch],
                      )
                  ]
              )
              .view(bsz, -1)
              .to(lprobs_list[i].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 MultichannelBeamSearch 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