colossalai.nn.layer.wrapper

class colossalai.nn.layer.wrapper.LambdaWrapper(func, layers_cfg=None)

Wrap a function to nn.Module, which takes a config of layers and can fully access them.

Parameters

• func (Callable) – User customed function.

• layers_cfg (dict, optional) – Config of layers, defaults to None.