code: def get_optim_state_fragment(self, key):
      if key in self.optim_fragment:
          return self.optim_fragment[key]
      else:
          raise ValueError(
              f"{key} not found in optimizer state fragment"
          )
  
file_path: H2O/h2o_flexgen/benchmark/third_party/DeepSpeed/deepspeed/utils/tensor_fragment.py
imports: import torch
  from dataclasses import dataclass
  from deepspeed import comm
label_desc: Use image node_id 3 for calling the tensor_fragment obj's underlying member method code with example usage: obj.get_optim_state_fragment(key) and returns: self