code: def forward(
      self,
      src_tokens: torch.Tensor,
      src_lengths: torch.Tensor,
      num_steps: int,
      dict_feat: Optional[
          Tuple[torch.Tensor, torch.Tensor, torch.Tensor]
      ] = None,
      contextual_token_embedding: Optional[torch.Tensor] = None,
  ):
      # Initialize incremental_states after every forward()
      self.decoder_ens.reset_incremental_states()
  
      # ARBABU: To do add an assert to make sure that batch size is 1
      decoder_ip = self.encoder_ens(
          src_tokens, src_lengths, dict_feat, contextual_token_embedding
      )
  
      (
          prev_token,
          prev_scores,
          prev_hypos_indices,
          attention_weights,
      ) = get_first_decoder_step_input(
          self.beam_size, self.target_dict_eos, src_lengths[0]
      )
  
      all_tokens_list = [prev_token]
      all_scores_list = [prev_scores]
      all_prev_indices_list = [prev_hypos_indices]
      all_attentions_list: List[torch.Tensor] = []
      if self.record_attention:
          all_attentions_list.append(attention_weights)
  
      for i in range(num_steps):
          (
              prev_token,
              prev_scores,
              prev_hypos_indices,
              attention_weights,
              decoder_ip,
          ) = self.decoder_ens(
              prev_token, prev_scores, i + 1, decoder_ip
          )
          all_tokens_list.append(prev_token)
          all_scores_list.append(prev_scores)
          all_prev_indices_list.append(prev_hypos_indices)
          if self.record_attention:
              all_attentions_list.append(attention_weights)
  
      all_tokens = torch.stack(all_tokens_list)
      all_scores = torch.stack(all_scores_list)
      all_prev_indices = torch.stack(all_prev_indices_list)
      if self.record_attention:
          all_attn_weights = torch.stack(all_attentions_list)
      else:
          all_attn_weights = torch.zeros(
              num_steps + 1, self.beam_size, src_tokens.size(1)
          )
  
      return all_tokens, all_scores, all_attn_weights, all_prev_indices
  
file_path: pytext/pytext/torchscript/seq2seq/beam_search.py
imports: from typing import List, Optional, Tuple
  import torch
  import torch.jit
  import torch.jit.quantized
  from pytext.common.constants import TORCH_VERSION
  from torch import nn
  from .decoder import DecoderBatchedStepEnsemble
  from .encoder import EncoderEnsemble
label_desc: Use image node_id 2 for calling the BeamSearch obj's underlying member method code with example usage: obj.forward(src_tokens, src_lengths, num_steps, dict_feat, contextual_token_embedding) and returns: all_tokens,  all_scores,  all_attn_weights,  all_prev_indices