From 96f6b843de0af7cdb6357e453f44dd55ea9f9c8d Mon Sep 17 00:00:00 2001 From: Chun-nien Chan Date: Tue, 17 Sep 2024 11:37:49 -0700 Subject: [PATCH] add _native_batch_norm_legit_functional decomp PiperOrigin-RevId: 675647257 --- ai_edge_torch/odml_torch/lowerings/registry.py | 1 + 1 file changed, 1 insertion(+) diff --git a/ai_edge_torch/odml_torch/lowerings/registry.py b/ai_edge_torch/odml_torch/lowerings/registry.py index e142a0af..b71b2a24 100644 --- a/ai_edge_torch/odml_torch/lowerings/registry.py +++ b/ai_edge_torch/odml_torch/lowerings/registry.py @@ -57,6 +57,7 @@ def register(self, op, lowering): torch._decomp.get_decompositions([ torch.ops.aten.upsample_nearest2d, torch.ops.aten._native_batch_norm_legit.no_stats, + torch.ops.aten._native_batch_norm_legit_functional, torch.ops.aten._adaptive_avg_pool2d, torch.ops.aten._adaptive_avg_pool3d, torch.ops.aten.grid_sampler_2d,