(
graph_module: torch.fx.GraphModule, node: torch.fx.Node, arg_index: int
)
| 37 | |
| 38 | |
| 39 | def _get_submodule( |
| 40 | graph_module: torch.fx.GraphModule, node: torch.fx.Node, arg_index: int |
| 41 | ) -> Tuple[str, torch.nn.Module, torch.fx.Node]: |
| 42 | submod_node = node.args[arg_index] |
| 43 | assert isinstance(submod_node, torch.fx.Node) |
| 44 | assert submod_node.op == "get_attr" |
| 45 | assert isinstance(submod_node.target, str) |
| 46 | submodule = graph_module.get_submodule(submod_node.target) |
| 47 | # pyre-ignore |
| 48 | return submod_node.target, submodule, node |
| 49 | |
| 50 | |
| 51 | def _get_control_flow_submodules( |