code: def assert_model_equal(
      m1: torch.nn.Module, m2: torch.nn.Module
  ) -> None:
      s1 = m1.state_dict()
      s2 = m2.state_dict()
  
      assert len(s1) == len(s2), f"len {len(s1)} vs {len(s2)}"
  
      for (n1, t1), (n2, t2) in zip(s1.items(), s2.items()):
          assert n1 == n2
          assert torch.equal(t1, t2), f"{n1} {t1} vs {t2}"
  
      for p1, p2 in zip(m1.parameters(), m2.parameters()):
          assert p1.requires_grad == p2.requires_grad
  
file_path: ColossalAI/tests/test_lazy/lazy_init_utils.py
imports: import random
  from copy import deepcopy
  from typing import Any, Callable, Optional, Tuple
  import numpy
  import torch
  from packaging import version
  from colossalai.device.device_mesh import DeviceMesh
  from colossalai.lazy.lazy_init import LazyInitContext, LazyTensor, _MyTensor
  from colossalai.tensor.d_tensor import to_global
  from colossalai.tensor.d_tensor.layout import Layout
  from tests.kit.model_zoo.registry import ModelAttribute
label_desc: Use image node_id 2 for calling a global function with example usage: assert_model_equal(m1, m2) without return types