diff --git a/core/lowering/passes/exception_elimination.cpp b/core/lowering/passes/exception_elimination.cpp index 63581af1ca..02fb773653 100644 --- a/core/lowering/passes/exception_elimination.cpp +++ b/core/lowering/passes/exception_elimination.cpp @@ -4,6 +4,7 @@ #include "torch/csrc/jit/passes/dead_code_elimination.h" #include "torch/csrc/jit/passes/guard_elimination.h" #include "torch/csrc/jit/passes/peephole.h" +#include "torch/csrc/jit/passes/remove_exceptions.h" #include "torch/csrc/jit/runtime/graph_executor.h" #include "core/util/prelude.h" @@ -21,6 +22,7 @@ struct ExceptionOrPassPatternElimination { void run() { findExceptionOrPassNodes(graph_->block()); + torch::jit::EliminateExceptions(graph_); torch::jit::EliminateDeadCode(graph_); LOG_GRAPH("Post exeception or pass elimination: " << *graph_); }