Add a comment regarding reference cycles

pytorch · Oct 24, 2024 · e8632b0 · e8632b0
1 parent 793f545
commit e8632b0
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/tensordict/_td.py b/tensordict/_td.py
@@ -978,6 +978,8 @@ def _cast_reduction(
                     agglomerate, keepdim=keepdim, dim=dim
                 )
 
+        # IMPORTANT: do not directly access batch_dims (or any other property)
+        # via self.batch_dims otherwise a reference cycle is introduced
         def proc_dim(dim, batch_dims, tuple_ok=True):
             if dim is None:
                 return dim