add a pass in ts_compile to prepare for jit.script (#899)

* add a pass in ts_compile to strip overloads and prepare for jit.script

Author: yushangdi

functorch/_src/compilers.py

@@ -20,6 +20,19 @@ def _canonicalize(fx_g):
     return fx_g
 
 
+def strip_overloads(gm):
+    """
+    Modifies the target of graph nodes in :attr:`gm` to strip overloads.
+
+    Args:
+        gm(fx.GraphModule): The input Fx graph module to be modified
+    """
+    for node in gm.graph.nodes:
+        if isinstance(node.target, torch._ops.OpOverload):
+            node.target = node.target.overloadpacket
+    gm.recompile()
+
+
 def ts_compile(fx_g: fx.GraphModule, _) -> Callable:
     """
     Compiles the :attr:`fx_g` with Torchscript compiler.
@@ -46,6 +59,8 @@ def ts_compile(fx_g: fx.GraphModule, _) -> Callable:
             new_kwargs[k] = v
         node.kwargs = new_kwargs
 
+    strip_overloads(fx_g)
+
     fx_g.graph.lint()
 
     fx_g.recompile()