Unverified Commit ec021923 authored by Patrick von Platen's avatar Patrick von Platen Committed by GitHub
Browse files

[Unet1d] correct docs (#2565)

parent 1598a579
......@@ -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.
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment