Layers

Poutyne provides utility layers that can be used with the Sequential module, ModuleList module and others.

class poutyne.Lambda(func)[source]

Applies a function to the input tensor.

Parameters:: func (Callable[[Tensor], Tensor]) – The function to apply.

Example

# Alternate version to the ``nn.Flatten`` module.
my_flatten = Lambda(lambda x: x.flatten(1))