[BugFix] Fix buffer identity in Params._apply (#1027)

pytorch · Oct 4, 2024 · 04faf40 · 04faf40 · github-actions · Oct 4, 2024
1 parent 362c072
commit 04faf40
Showing 1 changed file with 5 additions and 0 deletions.
diff --git a/tensordict/nn/params.py b/tensordict/nn/params.py
@@ -1164,7 +1164,12 @@ def _apply(self, fn, recurse=True):
         self._param_td._erase_cache()
         param_td = self._param_td
         self._param_td = param_td.copy()
+        # Keep a list of buffers to update .data only
+        bufs = dict(self._buffers)
         out: TensorDictBase = super()._apply(fn, recurse=recurse)
+        for key, val in bufs.items():
+            val.data = self._buffers[key].data
+            self._buffers[key] = val
         # Check device and shape
         cbs = out._check_batch_size(raise_exception=False)
         if not cbs: