class_docstr: Implementation of a model for inference composable with ragged batching.
class_docstr_tok: ["Implementation","of","a","model","for","inference","composable","with","ragged","batching","."]
docstr: Sets the state manager attribute. This is called by the inference engine after
  the model is fully initialized.
docstr_tok: ["Sets","the","state","manager","attribute",".","This","is","called","by","the","inference","engine","after","the","model","is","fully","initialized","."]
code: def set_state_manager(self, state_manager: DSStateManager):
      """
      Sets the state manager attribute. This is called by the inference engine after
      the model is fully initialized.
      """
      self.state_manager = state_manager
  
imports: from abc import ABC, abstractmethod
  from typing import Iterable, Optional, Tuple, Type
  import torch
  import deepspeed.comm
  from ..ragged import DSStateManager, RaggedBatchWrapper
  from ..ragged.manager_configs import KVCacheConfig
  from ..ragged import DSSequenceDescriptor
  from ..model_implementations.layer_container_base import LayerContainer
  from ..config_v2 import RaggedInferenceEngineConfig
  from .flat_model_helpers import ModelMetadata
label_desc: Use image node_id 4 for calling the DSInferenceModelBase obj's underlying member method code with example usage: obj.set_state_manager(state_manager) without return types