From ec021923d22b54235d0a2fecdfba6eade654583d Mon Sep 17 00:00:00 2001 From: Patrick von Platen Date: Mon, 6 Mar 2023 14:36:28 +0100 Subject: [PATCH] [Unet1d] correct docs (#2565) --- src/diffusers/models/unet_1d.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/diffusers/models/unet_1d.py b/src/diffusers/models/unet_1d.py index 19166a2b26..eada6ddf3a 100644 --- a/src/diffusers/models/unet_1d.py +++ b/src/diffusers/models/unet_1d.py @@ -195,7 +195,7 @@ class UNet1DModel(ModelMixin, ConfigMixin): ) -> Union[UNet1DOutput, Tuple]: r""" Args: - sample (`torch.FloatTensor`): `(batch_size, sample_size, num_channels)` noisy inputs tensor + sample (`torch.FloatTensor`): `(batch_size, num_channels, sample_size)` noisy inputs tensor timestep (`torch.FloatTensor` or `float` or `int): (batch) timesteps return_dict (`bool`, *optional*, defaults to `True`): Whether or not to return a [`~models.unet_1d.UNet1DOutput`] instead of a plain tuple.