code: def register_flop_count_impl(func):
      def wrapper(impl):
          FlopProfiler._custom_flop_count_impl[func] = impl
          return impl
  
      return wrapper
  
imports: from torch.fx import GraphModule
  from .passes import ShapeProp, graph_profile_pass, shape_prop_pass
  from .passes.graph_profile import FlopProfiler
label_desc: Use image node_id 1 for calling a global function with example usage: register_flop_count_impl(func) and returns: wrapper, impl