code: def check_dtype(model, expected_dtype):
      def find_dtype(module):
          for child in module.children():
              if isinstance(child, DeepSpeedTransformerInference):
                  return child.attention.attn_qkvw.dtype
              else:
                  found_dtype = find_dtype(child)
                  if found_dtype:
                      return found_dtype
  
      found_dtype = find_dtype(model)
      assert (
          found_dtype
      ), "Did not find DeepSpeedTransformerInference in model"
      assert (
          found_dtype == expected_dtype
      ), f"Expected transformer dtype {expected_dtype}, but found {found_dtype}"
  
file_path: H2O/h2o_flexgen/benchmark/third_party/DeepSpeed/tests/unit/inference/test_checkpoint_sharding.py
imports: import os
  import pytest
  import torch
  import deepspeed
  from deepspeed.model_implementations import DeepSpeedTransformerInference
  from transformers import AutoConfig, AutoModelForCausalLM
  from unit.common import DistributedTest, DistributedFixture
label_desc: Use image node_id 1 for calling a global function with example usage: check_dtype(model, expected_dtype) and returns: child, found_dtype