[BugFix] Fix mem leak when locking

ghstack-source-id: d6e44e1d9b9afc9903a0f45945c10a94dcf5a0ca
Pull Request resolved: #1188

tensordict/base.py

@@ -11144,12 +11144,12 @@ def __exit__(self, exc_type, exc_val, exc_tb):
             return False
         _last_op = self._last_op_queue.pop()
         if _last_op is not None:
-            last_op, (args, kwargs, out) = _last_op
+            last_op, (args, kwargs, out_wr) = _last_op
             # TODO: transpose, flatten etc. as decorator should lock the content to make sure that no key is
             #  added or deleted
             _inv_caller = LAST_OP_MAPS.get(last_op)
             if _inv_caller is not None:
-                return _inv_caller(self, args, kwargs, out)
+                return _inv_caller(self, args, kwargs, out_wr())
             else:
                 raise NotImplementedError(f"Unrecognised function {last_op}.")
         return self

tensordict/utils.py

@@ -18,6 +18,7 @@
 import threading
 import time
 import warnings
+import weakref
 from collections import defaultdict
 from collections.abc import KeysView
 from contextlib import nullcontext
@@ -1237,7 +1238,14 @@ def func_as_decorator(_self, *args, **kwargs):
                 _attr_post = getattr(_self, attr)
                 if out is not None:
                     if _attr_post is not _attr_pre:
-                        out._last_op = (func.__name__, (args, kwargs, _self))
+                        out._last_op = (
+                            func.__name__,
+                            (
+                                args,
+                                kwargs,
+                                weakref.ref(_self),
+                            ),
+                        )
                     else:
                         out._last_op = None
                 return out
@@ -1248,7 +1256,7 @@ def func_as_decorator(_self, *args, **kwargs):
             def func_as_decorator(_self, *args, **kwargs):
                 out = func(_self, *args, **kwargs)
                 if out is not None:
-                    out._last_op = (func.__name__, (args, kwargs, _self))
+                    out._last_op = (func.__name__, (args, kwargs, weakref.ref(_self)))
                 return out
 
         return func_as_decorator