Fix bug in graph partitioner

Summary:
X-link: pytorch/pytorch#125133

Title

Reviewed By: PaulZhang12

Differential Revision: D56688411

torchrec/distributed/tests/test_pt2.py

@@ -302,15 +302,18 @@ def test_sharded_quant_ebc_non_strict_export(self) -> None:
         kjt = kjt.to("meta")
         sharded_model(kjt.values(), kjt.lengths())
 
-        ep = torch.export.export(
+        from torch.export import _trace
+        ep_pre = _trace._export(
             sharded_model,
             (
                 kjt.values(),
                 kjt.lengths(),
             ),
             {},
+            pre_dispatch=True,
             strict=False,
         )
+        ep = ep_pre.run_decompositions()
 
         ep.module()(kjt.values(), kjt.lengths())
 
@@ -338,15 +341,18 @@ def test_sharded_quant_fpebc_non_strict_export(self) -> None:
 
         sharded_model(kjt.values(), kjt.lengths())
 
-        ep = torch.export.export(
+        from torch.export import _trace
+
+        ep = _trace._export(
             sharded_model,
             (
                 kjt.values(),
                 kjt.lengths(),
             ),
             {},
             strict=False,
-        )
+            pre_dispatch=True,
+        ).run_decompositions()
         ep.module()(kjt.values(), kjt.lengths())
 
         # PT2 IR autofunctionalizes mutation funcs (bounds_check_indices)