Skip to content

Commit e283b27

Browse files
lvliang-intelyiliu30
authored andcommitted
Fix NTM-One-Shot failed with KeyError (#210)
Signed-off-by: Lv, Liang1 <[email protected]> Signed-off-by: Lv, Liang1 <[email protected]> Signed-off-by: yiliu30 <[email protected]>
1 parent d7213db commit e283b27

File tree

1 file changed

+5
-5
lines changed

1 file changed

+5
-5
lines changed

neural_compressor/adaptor/tf_utils/graph_rewriter/generic/pre_optimize.py

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -146,16 +146,16 @@ def get_optimized_model(self, itex_mode=False):
146146

147147
self._tmp_graph_def = ConvertPlaceholderToConst(self._tmp_graph_def).do_transformation()
148148

149-
self._tmp_graph_def = RemoveTrainingNodesOptimizer(
150-
self._tmp_graph_def, protected_nodes=input_output_names).do_transformation()
151-
152149
self._tmp_graph_def = SwitchOptimizer(self._tmp_graph_def).do_transformation()
153150

151+
self._tmp_graph_def = GrapplerOptimizer(
152+
self._tmp_graph_def, input_output_names, self.optimization).do_transformation()
153+
154154
self._tmp_graph_def = StripUnusedNodesOptimizer(self._tmp_graph_def,
155155
input_node_names, output_node_names).do_transformation()
156156

157-
self._tmp_graph_def = GrapplerOptimizer(
158-
self._tmp_graph_def, input_output_names, self.optimization).do_transformation()
157+
self._tmp_graph_def = RemoveTrainingNodesOptimizer(
158+
self._tmp_graph_def, protected_nodes=input_output_names).do_transformation()
159159

160160
self._tmp_graph_def = SplitSharedInputOptimizer(self._tmp_graph_def).do_transformation()
161161

0 commit comments

Comments
 (0)