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

• colossalai.nn.layer.wrapper.lambda_wrapper
• colossalai.nn.layer.wrapper.pipeline_wrapper