[BugFix] Fix tensorclass update

ghstack-source-id: 3f50604
Pull Request resolved: #1255

Author: Vincent Moens

tensordict/nn/sequence.py

@@ -218,7 +218,10 @@ 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,
+                inplace=inplace,
             )
         elif len(modules) == 1 and isinstance(
             modules[0], collections.abc.MutableSequence
@@ -227,20 +230,27 @@ 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,
+                inplace=inplace,
             )
         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,
+                inplace=inplace,
             )
 
         self.inplace = inplace
@@ -628,6 +638,7 @@ def forward(
             )
         if tensordict_out is not None:
             result = tensordict_out
+            print('here! update')
             result.update(tensordict_exec, keys_to_update=self.out_keys)
         else:
             result = tensordict_exec

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