Fix the FUSE_PASSES_POST_OBSERVER call (#81)

Yinghai Lu · Wei Wei · commit 4a0ca3e9f32f · 2022-06-03T17:54:13.000-07:00
Summary: Pull Request resolved: https://github.com/pytorch/fx2trt/pull/81 The observer.observe calling signature is observe(module, inputs). But doing partial(observe, inputs) will reverse the order to observe(inputs, module). Fixing it by using an explicit lambda. Reviewed By: wushirong Differential Revision: D36638291 fbshipit-source-id: 64abc0952802d5b438e1013c5ff91a57442900d9
diff --git a/fx/passes/lower_pass_manager_builder.py b/fx/passes/lower_pass_manager_builder.py
@@ -96,7 +96,7 @@ def graph_optimization_pass(self) -> PassManager:
         for p in self.lower_setting.lower_basic_fuse_pass.passes:
             passes.append(wrapper(p, self._input))
         passes.append(
-            inplace_wrapper(partial(FUSE_PASSES_POST_OBSERVER.observe, self._input))
+            inplace_wrapper(lambda m: FUSE_PASSES_POST_OBSERVER.observe(m, self._input))
         )
 
         return PassManager.build_from_passlist(passes)

Original file line number	Diff line number	Diff line change
`@@ -96,7 +96,7 @@ def graph_optimization_pass(self) -> PassManager:`
`96`	`96`	`for p in self.lower_setting.lower_basic_fuse_pass.passes:`
`97`	`97`	`passes.append(wrapper(p, self._input))`
`98`	`98`	`passes.append(`
`99`		`- inplace_wrapper(partial(FUSE_PASSES_POST_OBSERVER.observe, self._input))`
	`99`	`+ inplace_wrapper(lambda m: FUSE_PASSES_POST_OBSERVER.observe(m, self._input))`
`100`	`100`	`)`
`101`	`101`
`102`	`102`	`return PassManager.build_from_passlist(passes)`