Avoid directly calling the map operator

Summary: Calling the higher order operator directly is not recommended since it skips the safety check.

Differential Revision: D74036464

Author: ydwu4

exir/tests/control_flow_models.py

@@ -6,6 +6,7 @@
 
 import torch
 from torch.nn import Module  # @manual
+from torch._higher_order_ops.map import map as torch_map
 
 
 class FTCondBasic(Module):
@@ -87,7 +88,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_map(f, xs, y) + xs
 
     def get_random_inputs(self):
         return torch.rand(2, 4), torch.rand(4)
@@ -101,7 +102,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_map(f, xs, y) + xs
 
     def get_upper_bound_inputs(self):
         return torch.rand(4, 4), torch.rand(4)