class_docstr: Inference model implementation for ragged batching for Llama-2 models.
class_docstr_tok: ["Inference","model","implementation","for","ragged","batching","for","Llama-2","models","."]
code: def forward(self, wrapped_batch: RaggedBatchWrapper) -> torch.Tensor:
      residual = self._forward_embed(wrapped_batch)
  
      residual, hidden_states = self.norm(
          residual,
          None,
          self._transformer[0].attn_norm_gamma,
          beta=None,
      )
  
      for layer_idx in range(self.num_layers):
          residual, hidden_states = self._forward_transformer_layer(
              layer_idx, residual, hidden_states, wrapped_batch
          )
  
      return self._forward_unembed(residual, wrapped_batch)
  
imports: from typing import Iterable, Optional, Tuple
  import torch
  import deepspeed.comm
  from ...allocator import empty_from
  from ...inference_utils import ActivationType, DtypeEnum
  from ..None import *
  from ...modules.configs import *
  from ...modules.interfaces import *
  from ...modules import heuristics
  from ...ragged import RaggedBatchWrapper
  from .container import Llama2NonTransformerContainer, Llama2TransformerContainer
label_desc: Use image node_id 17 for calling the Llama2InferenceModel obj's underlying member method code with example usage: obj.forward(wrapped_batch) and returns: self