code: def assert_forward_equal(
      m1: torch.nn.Module,
      m2: torch.nn.Module,
      data_gen_fn: Callable[[], dict],
      output_transform_fn: Callable[[Any], dict],
  ) -> None:
      data = data_gen_fn()
  
      m1.eval()
      m2.eval()
      # run forward
      with torch.no_grad():
          outputs1 = m1(**data)
          outputs2 = m2(**data)
  
      # compare output
      transformed_out1 = output_transform_fn(outputs1)
      transformed_out2 = output_transform_fn(outputs2)
  
      assert len(transformed_out1) == len(transformed_out2)
  
      for key, out1 in transformed_out1.items():
          out2 = transformed_out2[key]
          assert torch.allclose(
              out1, out2, atol=1e-5
          ), f"{m1.__class__.__name__} has inconsistent outputs, {out1} vs {out2}"
  
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 3 for calling a global function with example usage: assert_forward_equal(m1, m2, data_gen_fn, output_transform_fn) without return types