Override of `Tracer.create_proxy`. This override intercepts the recording of every operation and stores away the current traced module's qualified name in `node_to_originating_module`
(self, kind: str, target: torch.fx.node.Target, args: Tuple[Any, ...],
kwargs: Dict[str, Any], name: Optional[str] = None, type_expr: Optional[Any] = None)
| 48 | self.current_module_qualified_name = old_qualname |
| 49 | |
| 50 | def create_proxy(self, kind: str, target: torch.fx.node.Target, args: Tuple[Any, ...], |
| 51 | kwargs: Dict[str, Any], name: Optional[str] = None, type_expr: Optional[Any] = None): |
| 52 | """ |
| 53 | Override of `Tracer.create_proxy`. This override intercepts the recording |
| 54 | of every operation and stores away the current traced module's qualified |
| 55 | name in `node_to_originating_module` |
| 56 | """ |
| 57 | proxy = super().create_proxy(kind, target, args, kwargs, name, type_expr) |
| 58 | self.node_to_originating_module[proxy.node] = self.current_module_qualified_name |
| 59 | return proxy |
| 60 | |
| 61 | |
| 62 | # Testing: let's see how this works on a torchvision ResNet18 model |