code: def test_custom_formula():
      class Anything(nn.Module):
          def forward(self, x):
              return x
  
      def custom_formula(module, inputs, outputs):
          return FlopsResult(1.0, 2.0)
  
      net = Anything()
      register_flops_formula(Anything, custom_formula)
      profiler = FlopsParamsProfiler(net, torch.randn(1, 3))
      assert profiler.profile({}) == FlopsResult(1.0, 2.0)
  
      def new_custom_formula(module, inputs, outputs):
          return FlopsResult(2.0, 3.0)
  
      Anything._count_flops = new_custom_formula
      profiler = FlopsParamsProfiler(net, torch.randn(1, 3))
      assert profiler.profile({}) == FlopsResult(2.0, 3.0)
  
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 8 for calling a global function with example usage: test_custom_formula() and returns: FlopsResult, FlopsResult, x