amend

tensordict/_td.py

@@ -272,7 +272,6 @@ def to_module(self, module, return_swap: bool = True, swap_dest=None, memo=None)
             # is passed with to_module) but for the device it could be a problem.
             if swap_dest is None:
                 swap = TensorDict({}, batch_size=[])
-                swap.clear_device_()
             else:
                 swap = swap_dest
             memo[id(module)] = swap
@@ -324,6 +323,7 @@ def to_module(self, module, return_swap: bool = True, swap_dest=None, memo=None)
                 _swap[key] = local_out
         if return_swap:
             if isinstance(swap, TensorDict):
+                # this is very ad-hoc but faster than calling _set_str every time
                 swap._tensordict.update(_swap)
             else:
                 swap.update(_swap)

tensordict/base.py

@@ -3125,7 +3125,6 @@ def __exit__(self, exc_type, exc_val, exc_tb):
                 return self.lock_()
             if last_op == self.__class__.to_module.__name__:
                 if is_tensor_collection(out):
-                    # with out.unlock_():
                     return self.to_module(*args, **kwargs, swap_dest=out)
                 else:
                     raise RuntimeError(