code: def _check_data_types(keys, data, target_dtype):
      """Check that all the keys have the same target data type."""
      for key in keys:
          assert data[key].dtype == target_dtype, (
              "{} has data type {} which "
              "is different than {}".format(
                  key, data[key].dtype, target_dtype
              )
          )
  
imports: import torch
  from megatron.core.parallel_state import get_tensor_model_parallel_group, get_tensor_model_parallel_rank, get_tensor_model_parallel_src_rank
label_desc: Use image node_id 1 for calling a global function with example usage: _check_data_types(keys, data, target_dtype) without return types