[BugFix] Fix tensorclass update

ghstack-source-id: 3f50604b340b7a7cdc710dfaceedf563295b2911 Pull Request resolved: #1255
pytorch · Mar 6, 2025 · cb81b5e · cb81b5e
1 parent 55fab2a
commit cb81b5e
Show file tree

Hide file tree

Showing 3 changed files with 18 additions and 5 deletions.
diff --git a/tensordict/nn/common.py b/tensordict/nn/common.py
@@ -1089,9 +1089,9 @@ def __init__(
             )
 
         self.module = module
-        if inplace not in (True, False, "empty"):
+        if inplace not in (None, True, False, "empty"):
             raise ValueError(
-                f"The only accepted valued for inplace is `True`, `False`, or `'empty'`. Got inplace={inplace} "
+                f"The only accepted valued for inplace is `None`, `True`, `False`, or `'empty'`. Got inplace={inplace} "
                 "instead."
             )
         self.inplace = inplace

diff --git a/tensordict/nn/sequence.py b/tensordict/nn/sequence.py
@@ -218,7 +218,9 @@ def __init__(
                 **{key: val for key, val in _zip_strict(modules[0], modules_vals)}
             )
             super().__init__(
-                module=nn.ModuleDict(modules), in_keys=in_keys, out_keys=out_keys
+                module=nn.ModuleDict(modules),
+                in_keys=in_keys,
+                out_keys=out_keys,
             )
         elif len(modules) == 1 and isinstance(
             modules[0], collections.abc.MutableSequence
@@ -227,20 +229,25 @@ def __init__(
             in_keys, out_keys = self._compute_in_and_out_keys(modules)
             self._complete_out_keys = list(out_keys)
             super().__init__(
-                module=nn.ModuleList(modules), in_keys=in_keys, out_keys=out_keys
+                module=nn.ModuleList(modules),
+                in_keys=in_keys,
+                out_keys=out_keys,
             )
         elif len(modules) == 1 and isinstance(modules[0], dict):
             return self.__init__(
                 collections.OrderedDict(modules[0]),
                 partial_tolerant=partial_tolerant,
                 selected_out_keys=selected_out_keys,
+                inplace=inplace,
             )
         else:
             modules = self._convert_modules(modules)
             in_keys, out_keys = self._compute_in_and_out_keys(modules)
             self._complete_out_keys = list(out_keys)
             super().__init__(
-                module=nn.ModuleList(list(modules)), in_keys=in_keys, out_keys=out_keys
+                module=nn.ModuleList(list(modules)),
+                in_keys=in_keys,
+                out_keys=out_keys,
             )
 
         self.inplace = inplace

diff --git a/tensordict/tensorclass.py b/tensordict/tensorclass.py
@@ -1694,6 +1694,12 @@ def _update(
         update_batch_size=update_batch_size,
         ignore_lock=ignore_lock,
     )
+    # We also need to remove things from non_tensordict
+    if self._non_tensordict:
+        keys = set(self._tensordict.keys())
+        ntd = {k: val for k, val in self._non_tensordict.items() if k not in keys}
+        self._non_tensordict.clear()
+        self._non_tensordict.update(ntd)
     return self