diff --git a/exir/tests/control_flow_models.py b/exir/tests/control_flow_models.py
index 0815f8fbc5b..7e77de45430 100644
--- a/exir/tests/control_flow_models.py
+++ b/exir/tests/control_flow_models.py
@@ -87,7 +87,7 @@ def forward(self, xs, y):
         def f(x, y):
             return x + y
 
-        return torch.ops.higher_order.map(f, xs, y) + xs
+        return torch._higher_order_ops.map(f, xs, y) + xs
 
     def get_random_inputs(self):
         return torch.rand(2, 4), torch.rand(4)