Module bastionlab.torch.psg

Sub-modules

• bastionlab.torch.psg.convert
• bastionlab.torch.psg.nn

Functions

expand_weights(module: torch.nn.modules.module.Module, max_batch_size: int) ‑> None

Recursively converts the layers of a model to their expanded counterpart in [bastionlab.torch.psg.nn](nn.md).

Args: module: model whose weights must be expanded. max_batch_size: maximum size of the batches that will be processed by the model.