pytorch
2dd23ebf - Add support for multi output nodes in partial eval graph stitching (#66097)

Commit
3 years ago
Add support for multi output nodes in partial eval graph stitching (#66097) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/66097 Adding logic to generate runtime shapes for nodes with multi-outputs. It is generalizing existing flow of looking at a node, getting its shape graph, inlining it, and adding a mapping from the output to the new value in the stitched shape compute graph to loop over multiple outputs. Test Plan: Imported from OSS Reviewed By: navahgar Differential Revision: D31797468 Pulled By: eellison fbshipit-source-id: 2c182b71a46b36d33f23ad35b89790a4a5d4471c
Author
Elias Ellison
Parents
Loading