code: def step(
      self,
      step: int,
      lprobs,
      scores,
      prev_output_tokens: Optional[Tensor] = None,
      original_batch_idxs: Optional[Tensor] = None,
  ):
      n_channels = len(lprobs)
      bsz, beam_size, vocab_size = lprobs[0].size()
  
      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[i] = lprobs[i][:, ::beam_size, :].contiguous()
  
      probs = []
      top_indices = []
      for i in range(n_channels):
          if self.sampling_topp > 0:
              # only sample from the smallest set of words whose cumulative probability mass exceeds p
              probs_i, top_indices_i = self._sample_topp(lprobs[i])
          elif self.sampling_topk > 0:
              # only sample from top-k candidates
              lprobs[i], top_indices_i = lprobs[i].topk(
                  min(self.sampling_topk, lprobs[i].size(-1))
              )
              probs_i = lprobs[i].exp_()
          else:
              probs_i = lprobs[i].exp_()
  
              # dummy data to be consistent with true branch for type check
              top_indices_i = torch.empty(0).to(probs_i)
          probs.append(probs_i)
          top_indices.append(top_indices_i)
      # sample
      indices_buf = []
      for i in range(n_channels):
          if step == 0:
              indices_buf.append(
                  torch.multinomial(
                      probs[i].view(bsz, -1),
                      beam_size,
                      replacement=True,
                  ).view(bsz, beam_size)
              )
          else:
              indices_buf.append(
                  torch.multinomial(
                      probs[i].view(bsz * beam_size, -1),
                      1,
                      replacement=True,
                  ).view(bsz, beam_size)
              )
  
      if step == 0:
          for i in range(n_channels):
              # expand to beam size
              probs[i] = probs[i].expand(bsz, beam_size, -1)
  
      # gather scores
      scores_buf = []
      for i in range(n_channels):
          scores_buf.append(
              torch.gather(
                  probs[i], dim=2, index=indices_buf[i].unsqueeze(-1)
              )
          )
          scores_buf[i] = scores_buf[i].log_().view(bsz, -1)
  
      # remap indices if using top-k or top-P sampling
      if self.sampling_topk > 0 or self.sampling_topp > 0:
          for i in range(n_channels):
              indices_buf[i] = torch.gather(
                  top_indices[i].expand(bsz, beam_size, -1),
                  dim=2,
                  index=indices_buf[i].unsqueeze(-1),
              ).squeeze(2)
  
      if step == 0:
          beams_buf = indices_buf[0].new_zeros(bsz, beam_size)
      else:
          beams_buf = (
              torch.arange(0, beam_size)
              .to(indices_buf[0])
              .repeat(bsz, 1)
          )
          # make scores cumulative
          for i in range(n_channels):
              scores_buf[i].add_(
                  torch.gather(
                      scores[:, :, step - 1, i], dim=1, index=beams_buf
                  )
              )
      scores_buf = torch.stack(scores_buf, dim=-1)
      indices_buf = torch.stack(indices_buf, dim=-1)
  
      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 3 for calling the ContiguousMultichannelSampling 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