docstr: x:              input tensor
  hidden_state:   list of hidden_states, one for each resblock
                  if None, then becomes [None] * n_layer
  
  returns
  x:                      The output of the layers
  output_hidden_states:   A list of size self.resblocks with each hidden_state
docstr_tok: ["x",":","input","tensor","hidden_state",":","list","of","hidden_states",",","one","for","each","resblock","if","None",",","then","becomes","[","None","]","*","n_layer","returns","x",":","The","output","of","the","layers","output_hidden_states",":","A","list","of","size","self.resblocks","with","each","hidden_state"]
code: def forward(self, x, hidden_states=None):
      """
      x:              input tensor
      hidden_state:   list of hidden_states, one for each resblock
                      if None, then becomes [None] * n_layer
  
      returns
      x:                      The output of the layers
      output_hidden_states:   A list of size self.resblocks with each hidden_state
      """
      if hidden_states is None:
          hidden_states = [None] * len(self.resblocks)
      else:
          hidden_states = hidden_states
  
      assert len(hidden_states) == len(
          self.resblocks
      ), f"number of hidden states should match number of resblocks: {len(hidden_states)} != {len(self.resblocks)}"
  
      output_hidden_states = []
      # Blocks
      for l, hidden_state in zip(self.resblocks, hidden_states):
          x, output_hidden_state = l(x, hidden_state=hidden_state)
          output_hidden_states.append(output_hidden_state)
  
      return x, output_hidden_states
  
imports: import logging
  from abc import ABC, abstractmethod
  from dataclasses import dataclass
  from typing import List, Optional, Dict, TypeVar
  import numpy
  import torch
  import torch.nn
  import torch.nn.functional
  from torch.nn import LayerNorm
  from summarize_from_feedback.model_layout import ModelLayout
  from summarize_from_feedback.models.attention import Attention
  from summarize_from_feedback.models.ops import ACT_FNS, Conv1D
  from summarize_from_feedback.utils import exact_div
  from summarize_from_feedback.utils.dist_utils import create_model_parallel_comm, Comm
label_desc: Use image node_id 2 for calling the TransformerTorso obj's underlying member method code with example usage: obj.forward(x, hidden_states) and returns: x,  output_hidden_states