From 3e3b6dc14d74a5c2dd493d316c70afeedaa4850b Mon Sep 17 00:00:00 2001
From: vmoens <vincentmoens@gmail.com>
Date: Wed, 21 Feb 2024 10:06:16 -0800
Subject: [PATCH] init

---
 tensordict/_td.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/tensordict/_td.py b/tensordict/_td.py
index 24e909dab..0a825b88e 100644
--- a/tensordict/_td.py
+++ b/tensordict/_td.py
@@ -330,7 +330,8 @@ def _to_module(
         if not use_state_dict and isinstance(module, TensorDictBase):
             if return_swap:
                 swap = module.copy()
-                module.update(self)
+                with module.unlock_():
+                    module.update(self)
                 return swap
             else:
                 module.update(self)