Skip to content

nexuml.core.torch_adapter

nexuml.core.torch_adapter

Adapter wrapping plain torch.nn.Modules for TensorDict pipeline compatibility.

NnModuleLayer

Bases: LayerDefinition

Portable definition for an importable one-input/one-output PyTorch module.

TorchModuleAdapter

Bases: PipelineLayer

Wraps any torch.nn.Module to conform to PipelineLayer interface.

nn_module

nn_module(
    factory: Callable[P, Module],
    *args: args,
    **kwargs: kwargs,
) -> NnModuleLayer

Create a portable direct-module definition from an importable factory.

Returns:

Type Description
NnModuleLayer

Validated universal layer definition.

Raises:

Type Description
TypeError

If factory is a live module instance.