Make XLA tests export new IR compatible (#8354)

pytorch · Nov 5, 2024 · 81c4caa · 81c4caa
1 parent ca0a244
commit 81c4caa
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/test/stablehlo/test_export_fx_passes.py b/test/stablehlo/test_export_fx_passes.py
@@ -212,7 +212,8 @@ def forward(self, x, dim, weight, bias, eps):
     after_decomp_out = native_layer_norm_impl(*args)
     self.assertTrue(
         torch.allclose(before_decomp_out, after_decomp_out, atol=1e-6))
-    ep = export(m, args, dynamic_shapes=dynamic_shapes)
+    ep_training = export(m, args, dynamic_shapes=dynamic_shapes)
+    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)