class_docstr: Inference model implementation for ragged batching for Mistral 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_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 ...model_implementations import *
  from ...modules.configs import *
  from ...modules.interfaces import *
  from ...ragged import RaggedBatchWrapper
  from .container import MistralNonTransformerContainer, MistralTransformerContainer
label_desc: Use image node_id 17 for calling the MistralInferenceModel obj's underlying member method code with example usage: obj.forward(wrapped_batch) and returns: self