add _native_batch_norm_legit_functional decomp

PiperOrigin-RevId: 675647257
google-ai-edge · Sep 17, 2024 · 96f6b84 · 96f6b84
1 parent 0d481c4
commit 96f6b84
Showing 1 changed file with 1 addition and 0 deletions.
diff --git 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,