Update test_export_fx_passes.py

pytorch · Nov 4, 2024 · ac632b0 · ac632b0
1 parent ea31979
commit ac632b0
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/test/stablehlo/test_export_fx_passes.py b/test/stablehlo/test_export_fx_passes.py
@@ -213,7 +213,7 @@ def forward(self, x, dim, weight, bias, eps):
     self.assertTrue(
         torch.allclose(before_decomp_out, after_decomp_out, atol=1e-6))
     ep_training = export(m, args, dynamic_shapes=dynamic_shapes)
-    ep = ep.run_decompositions({})
+    ep = ep_training.run_decompositions({})
     decompose_dynamic_native_layer_norm(ep.graph_module)
     ep.graph_module.recompile()
     self.assertFalse('aten.native_layer_norm' in ep.graph_module.code)