code: def _real_shape_inference(model, *args, **kwargs):
      shapes = {}
  
      def _save_hook(module, input, output, name):
          try:
              input_shapes = (
                  tuple(i.shape for i in input)
                  if isinstance(input, tuple)
                  else input.shape
              )
              output_shapes = (
                  tuple(o.shape for o in output)
                  if isinstance(output, tuple)
                  else output.shape
              )
  
              shapes[name] = (input_shapes, output_shapes)
          except:
              pass
  
      hooks = []
      for name, module in model.named_modules():
          hooks.append(
              module.register_forward_hook(
                  partial(_save_hook, name=name)
              )
          )
  
      model(*args, **kwargs)
  
      for hook in hooks:
          hook.remove()
  
      return shapes
  
file_path: nni/test/algo/nas/profiler/test_shape.py
imports: import re
  from functools import partial
  import torch
  from nni.nas.hub.pytorch import *
  from nni.nas.profiler.pytorch.utils import *
label_desc: Use image node_id 1 for calling a global function with example usage: _real_shape_inference(model) and returns: shapes