azula.nn.utils

Miscellaneous neural network helpers.

Classes

skip_init

Creates a context in which weight initialization is skipped.

Functions

cpu_offload

Moves a module to a device and offloads it to CPU upon exit.

Descriptions

class azula.nn.utils.skip_init

Creates a context in which weight initialization is skipped.

Example

>>> with skip_init():
...    layer = nn.Linear(3, 5)
azula.nn.utils.cpu_offload(module, device)

Moves a module to a device and offloads it to CPU upon exit.

Parameters:

  • module (Module) – The module to offload.

  • device (device) – The target device.