move module attribute inplace update to leaf function in ManagedCollisionModule

Summary:
inplace update will cause unexpected module attribute mutation described in pytorch/pytorch#70449

by moving it to leaf function we guaranteed no side effect during fx tracing.

Differential Revision: D73448087

Author: seanx92

torchrec/modules/mc_modules.py

@@ -59,23 +59,27 @@ def _mcc_lazy_init(
     features: KeyedJaggedTensor,
     feature_names: List[str],
     features_order: List[int],
-    created_feature_order: bool,
-) -> Tuple[KeyedJaggedTensor, bool, List[int]]:  # features_order
+    created_feature_order: List[bool],
+) -> KeyedJaggedTensor:
     input_feature_names: List[str] = features.keys()
-    if not created_feature_order:
+    if not created_feature_order or not created_feature_order[0]:
         for f in feature_names:
             features_order.append(input_feature_names.index(f))
 
         if features_order == list(range(len(input_feature_names))):
-            features_order = torch.jit.annotate(List[int], [])
-        created_feature_order = True
+            features_order.clear()
+
+        if len(created_feature_order) > 0:
+            created_feature_order[0] = True
+        else:
+            created_feature_order.append(True)
 
     if len(features_order) > 0:
         features = features.permute(
             features_order,
         )
 
-    return (features, created_feature_order, features_order)
+    return features
 
 
 @torch.fx.wrap
@@ -298,6 +302,7 @@ class ManagedCollisionCollection(nn.Module):
 
     _table_to_features: Dict[str, List[str]]
     _features_order: List[int]
+    _created_feature_order: List[bool]  # use list for inplace update in leaf function
 
     def __init__(
         self,
@@ -338,7 +343,7 @@ def __init__(
         self._feature_names: List[str] = [
             feature for config in embedding_configs for feature in config.feature_names
         ]
-        self._created_feature_order = False
+        self._created_feature_order = [False]
         self._features_order = []
 
     def _create_feature_order(
@@ -360,11 +365,7 @@ def forward(
         self,
         features: KeyedJaggedTensor,
     ) -> KeyedJaggedTensor:
-        (
-            features,
-            self._created_feature_order,
-            self._features_order,
-        ) = _mcc_lazy_init(
+        features = _mcc_lazy_init(
             features,
             self._feature_names,
             self._features_order,