8000 [BugFix] Reference cycle in TensorDict._cast_reduction by egaznep · Pull Request #1056 · pytorch/tensordict · GitHub
[go: up one dir, main page]
More Web Proxy on the site http://driver.im/
Skip to content

[BugFix] Reference cycle in TensorDict._cast_reduction #1056

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 3 commits into from
Oct 24, 2024
Merged

[BugFix] Reference cycle in TensorDict._cast_reduction #1056

Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
e342189
Fixes #1055
egaznep Oct 23, 2024
793f545
Apply formatting changes
egaznep Oct 24, 2024
e8632b0
Add a comment regarding reference cycles
egaznep Oct 24, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
16 changes: 11 additions & 5 deletions tensordict/_td.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -978,24 +978,30 @@ def _cast_reduction(
agglomerate, keepdim=keepdim, dim=dim
)

def proc_dim(dim, tuple_ok=True):
# 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
if isinstance(dim, tuple):
if tuple_ok:
return tuple(_d for d in dim for _d in proc_dim(d, tuple_ok=False))
return tuple(
_d
for d in dim
for _d in proc_dim(d, batch_dims, tuple_ok=False)
)
return dim
if dim >= self.batch_dims or dim < -self.batch_dims:
if dim >= batch_dims or dim < -batch_dims:
raise RuntimeError(
"dim must be greater than or equal to -tensordict.batch_dims and "
"smaller than tensordict.batch_dims"
)
if dim < 0:
return (self.batch_dims + dim,)
return (batch_dims + dim,)
return (dim,)

if dim is not NO_DEFAULT:
dim = proc_dim(dim, tuple_ok=tuple_ok)
dim = proc_dim(dim, self.batch_dims, tuple_ok=tuple_ok)
if not tuple_ok:
dim = dim[0]
if dim is not NO_DEFAULT or keepdim:
Expand Down
Loading
0