Buffer#

API#

class warp_nn.modules.buffer.Buffer(data: array, *, requires_grad: bool = False)[source]#

Bases: object

Class representing a non-learnable state array (e.g. running statistics).

Unlike parameters, buffers are not optimized, but they are included in the module’s state dictionary and moved along with the module.

Note

All buffers in a Module are of the Buffer type.

Parameters:

  • data – The underlying data container (array) of the buffer.

  • requires_grad – Whether the buffer requires gradients.

to(
device: Device,
) → Buffer[source]#

Move the buffer to the specified device.

Parameters:

device – The device to move the buffer to.

Returns:

The buffer itself.

property data: array[source]#

The underlying data container (array) of the buffer.

property device: Device[source]#

Device on which the buffer is allocated.

property dtype: type[source]#

Data type of the buffer.

property requires_grad: bool[source]#

Whether the buffer requires gradients.

property shape: tuple[int, ...][source]#

Shape of the buffer.