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,