code: def check_state_dict_equal_pytree(
      d1: OrderedDict, d2: OrderedDict, ignore_device: bool = True
  ):
      flat_d1, _ = tree_flatten(d1)
      flat_d2, _ = tree_flatten(d2)
      assert len(flat_d1) == len(flat_d2)
      for v1, v2 in zip(flat_d1, flat_d2):
          if isinstance(v1, torch.Tensor):
              assert isinstance(v2, torch.Tensor)
              if not ignore_device:
                  v1 = v1.to("cpu")
                  v2 = v2.to("cpu")
              assert_close_loose(v1, v2)
          else:
              assert v1 == v2, f"{v1} not equals to {v2}"
  
file_path: ColossalAI/colossalai/testing/comparison.py
imports: from typing import Any, List, OrderedDict
  import torch
  import torch.distributed
  from torch import Tensor
  from torch.distributed import ProcessGroup
  from torch.testing import assert_close
  from torch.utils._pytree import tree_flatten
label_desc: Use image node_id 6 for calling a global function with example usage: check_state_dict_equal_pytree(d1, d2, ignore_device) without return types