Fix BatchNorm tests by passing relevant input buffers.

zjgarvey · zjgarvey · commit 236c07d6c95e · 2025-10-20T10:56:10.000-07:00
Signed-off-by: zjgarvey &lt;zjgarvey@gmail.com&gt;
diff --git a/projects/e2e/torch_mlir_e2e_test/configs/fx_importer_backend.py b/projects/e2e/torch_mlir_e2e_test/configs/fx_importer_backend.py
@@ -149,9 +149,12 @@ def _export_run(self, artifact: torch.nn.Module, trace: Trace) -> Trace:
             )
             module = self._backend.compile(module)
             backend_module = self._backend.load(module)
+            input_buffers = prog.graph_signature.inputs_to_buffers.values()
             params = {
                 # **dict(artifact.named_parameters(remove_duplicate=False)),
-                **dict(artifact.named_buffers(remove_duplicate=False)),
+                name: value
+                for (name, value) in artifact.named_buffers(remove_duplicate=False)
+                if name in input_buffers
             }
             params_flat, params_spec = pytree.tree_flatten(params)
             params_flat = list(params_flat)

Original file line number	Diff line number	Diff line change
`@@ -149,9 +149,12 @@ def _export_run(self, artifact: torch.nn.Module, trace: Trace) -> Trace:`
`149`	`149`	`)`
`150`	`150`	`module = self._backend.compile(module)`
`151`	`151`	`backend_module = self._backend.load(module)`
	`152`	`+ input_buffers = prog.graph_signature.inputs_to_buffers.values()`
`152`	`153`	`params = {`
`153`	`154`	`# **dict(artifact.named_parameters(remove_duplicate=False)),`
`154`		`- **dict(artifact.named_buffers(remove_duplicate=False)),`
	`155`	`+ name: value`
	`156`	`+ for (name, value) in artifact.named_buffers(remove_duplicate=False)`
	`157`	`+ if name in input_buffers`
`155`	`158`	`}`
`156`	`159`	`params_flat, params_spec = pytree.tree_flatten(params)`
`157`	`160`	`params_flat = list(params_flat)`