code: def test_e2e_simple(model):
      model_space = MODELS[model]()
      if model == "custom_op":
          with pytest.raises(
              RuntimeError, match="Shape inference failed"
          ):
              FlopsParamsProfiler(
                  model_space, torch.randn(1, 1, 28, 28)
              )
          return
      if model == "multihead_attention":
          profiler = FlopsParamsProfiler(
              model_space,
              (
                  (
                      torch.randn(2, 128),
                      torch.randn(2, 128),
                      torch.randn(2, 128),
                  ),
              ),
          )
      else:
          profiler = FlopsParamsProfiler(
              model_space, torch.randn(1, 1, 28, 28)
          )
      sample = {}
      for _ in model_space.grid(memo=sample):
          # TODO: Check the result
          print(sample, profiler.flops_result.freeze(sample))
  
file_path: nni/test/ut/nas/profiler/test_flops.py
imports: import pytest
  import torch
  from torch import nn
  import nni
  import nni.nas.nn.pytorch
  from nni.mutable import Categorical
  from nni.nas.profiler.pytorch.flops import FlopsParamsProfiler, FlopsProfiler, NumParamsProfiler, FlopsResult, register_flops_formula
  from ut.nas.nn.models import MODELS
label_desc: Use image node_id 9 for calling a global function with example usage: test_e2e_simple(model) without return types