model.tools.generator.Storage

class Storage(auto_reset: bool = True)

Bases: Module

Stores the forward pass values

Parameters:

auto_reset (bool, optional) – If true, automatically clears the storage after the last call of the model generator. Set to False if you want the value to persist after a forward pass of the model. Defaults to True.

Methods

add_input	Adds information about the input tensor that the storage will expect
add_requests	Adds one to the receive counter threshold
forward	Store the input tensor and returns it back
get	Returns a list of saved tensors
reset	Resets storage
shape	Returns a list of channel values for the expected data

Attributes

training

add_input(channels: int) → None

Adds information about the input tensor that the storage will expect

This is only necessary so that the generator can calculate the parameters for subsequent layers.

Parameters:

channels (int) – Number of input tensor channels

add_requests() → None

Adds one to the receive counter threshold

A receive is considered a call to the get method. If the threshold is not zero, then the storage will be automatically released when the threshold is reached.

forward(X: Tensor) → Tensor

Store the input tensor and returns it back

get() → List[Tensor]

Returns a list of saved tensors

reset() → None

Resets storage

shape() → List[int]

Returns a list of channel values for the expected data

The method does not analyze the current data, but relies on the data received via the add_input method.