8000 [BUG] chunk method not work for nested tensor in TensorDict · Issue #1338 · pytorch/tensordict · GitHub
[go: up one dir, main page]
More Web Proxy on the site http://driver.im/
Skip to content
8000

[BUG] chunk method not work for nested tensor in TensorDict #1338

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

Jump to bottom
Closed
3 tasks
vermouth1992 opened this issue Jun 7, 2025 · 0 comments · Fixed by #1339
Closed
3 tasks

[BUG] chunk method not work for nested tensor in TensorDict #1338

vermouth1992 opened this issue Jun 7, 2025 · 0 comments · Fixed by #1339
Assignees
@vmoens
Labels
bug Something isn't working

Comments

@vermouth1992
Copy link
vermouth1992 commented Jun 7, 2025 •
edited
Loading

Describe the bug

chunk method not work for nested tensor in TensorDict, but it works fine in raw nested tensor

To Reproduce

import torch
from tensordict import tensorclass, TensorDict

vocab_size = 128
a_seq = torch.randint(low=0, high=vocab_size, size=(10,))
b_seq = torch.randint(low=0, high=vocab_size, size=(20,))
rmpad_seq = torch.nested.as_nested_tensor([a_seq, b_seq], layout=torch.jagged)
print(rmpad_seq.chunk(2))  # works fine
rmpad_batch = TensorDict.from_dict({'input_ids': rmpad_seq}, batch_size=[2]) 
rmpad_batch.chunk(2) # doesn't work

Expected behavior

A clear and concise description of what you expected to happen.

Screenshots

If applicable, add screenshots to help explain your problem.

System info

Pytorch 2.6 + Tensordict 0.8.3

import tensordict, numpy, sys, torch
print(tensordict.__version__, numpy.__version__, sys.version, sys.platform, torch.__version__)

Additional context

Add any other context about the problem here.

Reason and Possible fixes

If you know or suspect the reason for this bug, paste the code lines and suggest modifications.

Checklist

__

  • I have checked that there is no similar issue in the repo (required)
  • I have read the documentation (required)
  • I have provided a minimal working example to reproduce the bug (required)
@vermouth1992 vermouth1992 added the bug Something isn't working label Jun 7, 2025
@vmoens vmoens linked a pull request Jun 7, 2025 that will close this issue
[BugFix] Fix chunk of NJTs #1339
Merged
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Assignees

@vmoens vmoens

Labels
bug Something isn't working
Projects
None yet
Milestone
No milestone
Development

Successfully merging a pull request may close this issue.

[BugFix] Fix chunk of NJTs pytorch/tensordict
2 participants
@vermouth1992 @vmoens
0