class_docstr: Auxiliary module that wraps computation for a single decode step for a SockeyeModel.
  End-to-end traceable. Return values are put into a flat list to avoid return type constraints
  for traced modules.
class_docstr_tok: ["Auxiliary","module","that","wraps","computation","for","a","single","decode","step","for","a","SockeyeModel",".","End-to-end","traceable",".","Return","values","are","put","into","a","flat","list","to","avoid","return","type","constraints","for","traced","modules","."]
code: def forward(
      self,
      step_input,
      states: List[pt.Tensor],
      vocab_slice_ids: Optional[pt.Tensor] = None,
  ) -> List[pt.Tensor]:
      target_embed = self.embedding_target(step_input.unsqueeze(1))
      decoder_out, new_states = self.decoder(target_embed, states)
      decoder_out = decoder_out.squeeze(1)
  
      # step_output: (batch_size, target_vocab_size or vocab_slice_ids)
      step_output = self.output_layer(decoder_out, vocab_slice_ids)
  
      # return values are collected in a flat list due to constraints in mixed return types in traced modules
      # (can only by tensors, or lists of tensors or dicts of tensors, but no mix of them).
      outputs = [step_output, decoder_out]
      if self.has_target_factors:
          outputs += [
              fol(decoder_out) for fol in self.factor_output_layers
          ]
      outputs += new_states
      return outputs
  
imports: import copy
  import logging
  import os
  import time
  from dataclasses import dataclass
  from functools import lru_cache
  from typing import cast, Dict, List, Optional, Tuple, Union
  import numpy
  import torch
  from sockeye import __version__
  from .None import constants
  from .None import data_io
  from .None import decoder
  from .None import encoder
  from .None import layers
  from .None import transformer
  from .None import utils
  from .None import vocab
  from .config import Config
  from .encoder import FactorConfig
  from .layers import LengthRatioConfig
  from .None import nvs
  from sockeye.knn import KNNConfig
label_desc: Use image node_id 2 for calling the _DecodeStep obj's underlying member method code with example usage: obj.forward(step_input, states, vocab_slice_ids) and returns: outputs