code: def get_shape(i):
      if isinstance(i, torch.Tensor):
          return i.shape
      elif hasattr(i, "weight"):
          return i.weight().shape
      else:
          raise ValueError(f"Unknown type {type(i)}")
  
imports: import os
  from collections import defaultdict
  from numbers import Number
  from typing import Any, List
  import torch
  from torch.utils._python_dispatch import TorchDispatchMode
  from torch.utils._pytree import tree_map
  from torchvision.models._api import Weights
label_desc: Use image node_id 1 for calling a global function with example usage: get_shape(i) and returns: i, i