freenode
AI & ML

PyTorch proposes BufferDict for registered buffer collections

A new torch.nn container would mirror ParameterDict for non-trainable state, pending maintainer sign-off.

PyTorch contributors have proposed torch.nn.BufferDict, a public dict-style container that registers each entry as a module buffer and lets callers control persistence per key.

The addition fills a gap beside the existing ParameterDict and ModuleDict helpers: models often keep tensors that are not parameters (running statistics, caches, masks) yet still need to ride along in state_dict, device moves, and export. BufferDict is meant to hold those tensors with Dynamo and export support, equality helpers for tests, and documentation.

An automated review of the change found nothing blocking merge on correctness grounds, but recommended discussion before landing a new public torch.nn surface. It also flagged edge cases when a key is under parametrization: assigning a Buffer with persistent=False can silently lose that flag, and deletion can recompute parametrized values, briefly register parameters on a container that claims not to hold them, or fail if the parametrization forward raises. Those persistence and cleanup behaviors need a clear policy or fix before the API is considered stable.

The proposal therefore waits on neural-network maintainer review rather than routine merge.