(doc/fluxion/pixelshuffle) add/convert docstrings to mkdocstrings format

This commit is contained in:
Laurent 2024-02-01 21:47:27 +00:00 committed by Laureηt
parent 18682f8155
commit be75f68893

View file

@ -4,5 +4,16 @@ from refiners.fluxion.layers.module import Module
class PixelUnshuffle(_PixelUnshuffle, Module): class PixelUnshuffle(_PixelUnshuffle, Module):
"""Pixel Unshuffle layer.
This layer wraps [`torch.nn.PixelUnshuffle`][torch.nn.PixelUnshuffle].
Receives:
(Float[Tensor, "batch in_channels in_height in_width"]):
Returns:
(Float[Tensor, "batch out_channels out_height out_width"]):
"""
def __init__(self, downscale_factor: int): def __init__(self, downscale_factor: int):
_PixelUnshuffle.__init__(self, downscale_factor=downscale_factor) _PixelUnshuffle.__init__(self, downscale_factor=downscale_factor)