fix internal error with reorder submodules (#146181)

avikchaudhuri · pytorchmergebot · commit cde5ddfd1460 · 2025-02-01T00:30:42.000Z
Test Plan: hard to isolate as small repro Differential Revision: D68963033 Pull Request resolved: #146181 Approved by: https://github.com/angelayi
diff --git a/torch/export/unflatten.py b/torch/export/unflatten.py
@@ -1389,7 +1389,7 @@ def _reorder_submodules(
         if child is None:
             continue
         fqn = prefix + name
-        _reorder_submodules(child, fqn_order, prefix=fqn + ".")
+        _reorder_submodules(child, fqn_order, prefix=fqn.split("@")[0] + ".")
         delattr(parent, name)
         children.append((fqn_order[fqn], name, child))
     children.sort(key=operator.itemgetter(0))
