diff --git a/imperative/python/megengine/functional/nn.py b/imperative/python/megengine/functional/nn.py index e862c760..fdc9eb7e 100644 --- a/imperative/python/megengine/functional/nn.py +++ b/imperative/python/megengine/functional/nn.py @@ -1595,24 +1595,16 @@ def sliding_window_transpose( def pad( src: Tensor, pad_witdth: Tuple[Tuple[int, int], ...], - mode: str = "CONSTANT", + mode: str = "constant", constant_value: float = 0.0, ) -> Tensor: """ - pad + Pad is python warpper for padding opr in megbrain, can padding in random one of the max 7 dimensions. + Supported constant, edge(replicate) and reflect mode, constatnt is the default mode. """ p_offsets = [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0] - assert mode in [ - "constant", - "CONSTANT", - "edge", - "EDGE", - "replicate", - "REPLICATE", - "reflect", - "REFLECT", - ] + assert mode.lower() in ["constant", "edge", "replicate", "reflect"] if mode.lower() == "edge": mode = "replicate" diff --git a/imperative/python/megengine/module/padding.py b/imperative/python/megengine/module/padding.py index 0dcc3714..433ecd32 100644 --- a/imperative/python/megengine/module/padding.py +++ b/imperative/python/megengine/module/padding.py @@ -5,10 +5,15 @@ from .module import Module class Pad(Module): + """ + Pad is python warpper for padding opr in megbrain, can padding in random one of the max 7 dimensions. + Supported constant, edge(replicate) and reflect mode, constatnt is the default mode. + """ + def __init__( self, pad_witdth: Tuple[Tuple[int, int], ...], - mode: str = "CONSTANT", + mode: str = "constant", constant_val: float = 0.0, ): super().__init__()