From ac632b06516c14ec7f30480d4799c32fb25d7045 Mon Sep 17 00:00:00 2001 From: Tugsbayasgalan Manlaibaatar Date: Mon, 4 Nov 2024 14:34:43 -0700 Subject: [PATCH] Update test_export_fx_passes.py --- test/stablehlo/test_export_fx_passes.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/test/stablehlo/test_export_fx_passes.py b/test/stablehlo/test_export_fx_passes.py index 2d26a1a1334..580b756e177 100644 --- 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)