code: def check_equal(tensor_a, tensor_b, atol=1e-06):
      assert (
          torch.allclose(tensor_a, tensor_b, rtol=0, atol=atol) is True
      )
  
file_path: ColossalAI/tests/test_moe/test_kernel.py
imports: import pytest
  import torch
  import torch.distributed
  import colossalai
  from colossalai.moe import SparseMLP
  from colossalai.moe.manager import MOE_MANAGER
  from colossalai.testing import rerun_if_address_is_in_use, spawn
  from colossalai.utils import get_current_device
label_desc: Use image node_id 1 for calling a global function with example usage: check_equal(tensor_a, tensor_b, atol) without return types