colossalai.nn.layer.wrapper.lambda_wrapper

class colossalai.nn.layer.wrapper.lambda_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