code: def test_flop_count_module(m):
      x = torch.rand(2, 3, 224, 224)
      with MetaTensorMode():  # save time for testing
          module = m()
      rs_fwd, rs_bwd = flop_count(module, x, verbose=True)
      assert rs_fwd > 0, f"fwd flop count of {m.__name__} is {rs_fwd}"
      assert rs_bwd > 0, f"bwd flop count of {m.__name__} is {rs_bwd}"
  
file_path: ColossalAI/tests/test_analyzer/test_subclasses/test_flop_tensor.py
imports: import pytest
  import torch
  import torch.nn.functional
  import torchvision.models
  from packaging import version
  from tests.test_analyzer.test_fx.zoo import tm_models, tmm_models
label_desc: Use image node_id 1 for calling a global function with example usage: test_flop_count_module(m) without return types