Searched refs:sequential_split (Results 1 – 4 of 4) sorted by relevance
| /aosp_15_r20/external/pytorch/torch/_export/passes/ |
| H A D | replace_set_grad_with_hop_pass.py | 6 from ..utils import node_inline_, nodes_filter, nodes_first, nodes_map, sequential_split
88 new_gm = sequential_split(gm, _is_set_grad_enabled_node)
|
| H A D | replace_autocast_with_hop_pass.py | 7 from ..utils import node_inline_, nodes_filter, nodes_first, sequential_split
135 return sequential_split(gm, node_call_back)
|
| /aosp_15_r20/external/pytorch/test/export/ |
| H A D | test_passes.py | 35 sequential_split,
869 new_gm = sequential_split(gm, _is_set_grad_enabled_node)
879 new_gm = sequential_split(gm, _is_set_grad_enabled_node)
1087 new_gm = sequential_split(gm, _is_set_grad_enabled_node)
|
| /aosp_15_r20/external/pytorch/torch/_export/ |
| H A D | utils.py | 485 def sequential_split(gm: torch.fx.GraphModule, node_call_back) -> torch.fx.GraphModule: function
|