Make dict update fx compatible (#2766)

Summary:
Pull Request resolved: #2766

Dict's `update` method cannot be supported by *fx*:
```
NotImplementedError: 'immutable_dict' object does not support mutation. If you are attempting to modify the kwargs or args of a torch.fx.Node object,
```

It needs to be handled for every key-value pair.

Not 100% sure why this issue was not surfaced before. But likely it is because we have never published a model with more than one MPZCH-powered embedding table, so the code path was not hit before.

Reviewed By: beekbin

Differential Revision: D70227971

fbshipit-source-id: 71ef081eb68fa7c3d91d1e3431bd88bc6673285c

Author: Ziliang Zhao

torchrec/modules/mc_modules.py

@@ -40,11 +40,13 @@ def apply_mc_method_to_jt_dict(
 def _update(
     base: Optional[Dict[str, JaggedTensor]], delta: Dict[str, JaggedTensor]
 ) -> Dict[str, JaggedTensor]:
-    if base is None:
-        base = delta
-    else:
-        base.update(delta)
-    return base
+    res: Dict[str, JaggedTensor] = {}
+    if base is not None:
+        for k, v in base.items():
+            res[k] = v
+    for k, v in delta.items():
+        res[k] = v
+    return res
 
 
 @torch.fx.wrap