neural_lam.utils.buffer_list

neural_lam.utils.buffer_list#

Module wrapper exposing a list of tensors as registered buffers.

Classes#

BufferList

A list of torch buffer tensors that sit together as a Module with no

Module Contents#

class neural_lam.utils.buffer_list.BufferList(buffer_tensors: list[torch.Tensor], persistent: bool = True)[source]#

Bases: torch.nn.Module

A list of torch buffer tensors that sit together as a Module with no parameters and only buffers.

This should be replaced by a native torch BufferList once implemented. See: pytorch/pytorch#37386

Register a collection of tensors as buffers inside a module.

Parameters:

  • buffer_tensors (Sequence[torch.Tensor]) – Buffers to register in the order they should be indexed.

  • persistent (bool, optional) – If True, buffers are saved in checkpoints. Default True.

n_buffers#