Skip to content

[BugFix] Use torch.nn.Buffer instead of the tensordict placeholder - #4541

Open
ahlag wants to merge 1 commit into
pytorch:mainfrom
ahlag:fix-4535-buffer
Open

ahlag wants to merge 1 commit into
pytorch:mainfrom
ahlag:fix-4535-buffer

Conversation

@ahlag

@ahlag ahlag commented Oct 9, 2026 •

Copy link
Copy Markdown

Description

On torch>=2.5, tensordict.utils.Buffer is an empty placeholder class, so isinstance(x, Buffer) is always False and the buffer branches in _cast, _map_weight and _to_module_compatible_tensor never run: a torch.nn.Buffer comes back as a plain tensor.

  • torchrl/_utils.py now imports torch.nn.Buffer, falling back to the tensordict copy on torch<2.5 (same try/except ImportError pattern as is_compiling just above it). collectors/utils.py, weight_update/weight_sync_schemes.py and objectives/common.py import Buffer from there.
  • _make_target_param checked isinstance(x, nn.Parameter) after taking x.data, which is never a Parameter. It now checks the input. This doesn't change what users see, since TensorDictParams already registers the target params as buffers, but the branch now does what the comment above it says.
  • test_weight_helpers_keep_buffers in test/test_collectors.py covers the three helpers (fails on main, passes with the fix; skipped on torch<2.5, where nn.Buffer doesn't exist).

Tested on macOS (CPU), torch 2.14.1, tensordict 0.14.3: the new test plus 129 related collector/loss tests, and test/objectives/test_loss_module.py, test_dqn.py, test_ddpg.py (1491 passed). ufmt and flake8 clean at the pinned pre-commit versions. I couldn't run torch<2.5 locally, so could a maintainer add the ci/olddeps label?

I used an AI assistant (Claude Code) to help write this change; I've reviewed and tested it.

Motivation and Context

close #4535

  • I have raised an issue to propose this change (required for new features and bug fixes)

Types of changes

  • Bug fix (non-breaking change which fixes an issue)

Checklist

  • I have read the CONTRIBUTION guide (required)
  • My change requires a change to the documentation.
  • I have updated the tests accordingly (required for a bug fix or a new feature).
  • I have updated the documentation accordingly.

@pytorch-bot

pytorch-bot Bot commented Oct 9, 2026

Copy link
Copy Markdown

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/rl/4541

Note: Links to docs will display an error until the docs builds have been completed.

This comment was automatically generated by Dr. CI and updates every 15 minutes.

@meta-cla meta-cla Bot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label Oct 9, 2026
On torch>=2.5, tensordict.utils.Buffer is an empty placeholder class, so
isinstance checks against it are always False and the buffer branches in
_cast, _map_weight and _to_module_compatible_tensor never run. Import
torch.nn.Buffer from torchrl._utils, falling back to the tensordict copy
on older torch.

_make_target_param also checked isinstance(x, nn.Parameter) after taking
x.data, which is never a Parameter; check the input instead.

Fixes pytorch#4535
@peterdsharpe

Copy link
Copy Markdown
Contributor

This is the right fix, but don't merge this until torchRL pins a torch>=2.5

@bsprenger
bsprenger self-requested a review October 10, 2026 17:16

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

BugFix CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. Collectors Integrations/torch_geometric Integrations Objectives WeightUpdate

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[BUG] torchrl uses tensordict.utils.Buffer as torch.nn.Buffer, but it is an empty class on torch>=2.5

2 participants