From f90153cda30d3f9aaa74fb7af7566db51d380457 Mon Sep 17 00:00:00 2001 From: Jordan Fix Date: Wed, 8 Sep 2021 15:30:28 -0700 Subject: [PATCH] [acc_normalizer] Improve error when kwarg normalization fails (#64408) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/64408 att Test Plan: NFC Reviewed By: protonu Differential Revision: D30716392 fbshipit-source-id: e1c3bb1afcd5363a9d502549d8a46b90226be40c --- torch/fx/experimental/fx_acc/acc_normalizer.py | 13 ++++++++++--- 1 file changed, 10 insertions(+), 3 deletions(-) diff --git a/torch/fx/experimental/fx_acc/acc_normalizer.py b/torch/fx/experimental/fx_acc/acc_normalizer.py index a5d116a..66e83a5 100644 --- a/torch/fx/experimental/fx_acc/acc_normalizer.py +++ b/torch/fx/experimental/fx_acc/acc_normalizer.py @@ -372,9 +372,16 @@ def normalize(mod: torch.fx.GraphModule, expect_nodes_have_shapes: bool = False) normalized_kwargs = node.kwargs else: normalized_args = () - normalized_kwargs = get_normalized_kwargs( - node, normalization_info.arg_replacement_tuples - ) + try: + normalized_kwargs = get_normalized_kwargs( + node, normalization_info.arg_replacement_tuples + ) + except Exception: + print( + f"Error during kwarg normalization for: {node.format_node()}; " + f"arg_replacement_tuples={normalization_info.arg_replacement_tuples}" + ) + raise if ( normalization_info.needs_shapes_for_normalization -- 2.7.4