"vscode:/vscode.git/clone" did not exist on "cc85a2e522055b11e2ce2515f078d3978b35d4b4"
Unverified Commit f1b9ee7e authored by Kashif Rasul's avatar Kashif Rasul Committed by GitHub
Browse files

[Docs] fix docstring for issue #709 (#710)

fix docstring

fixes #709
parent 4ff4d4db
......@@ -170,7 +170,7 @@ class UNet2DModel(ModelMixin, ConfigMixin):
timestep: Union[torch.Tensor, float, int],
return_dict: bool = True,
) -> Union[UNet2DOutput, Tuple]:
"""r
r"""
Args:
sample (`torch.FloatTensor`): (batch, channel, height, width) noisy inputs tensor
timestep (`torch.FloatTensor` or `float` or `int): (batch) timesteps
......
......@@ -215,7 +215,7 @@ class FlaxUNet2DConditionModel(nn.Module, FlaxModelMixin, ConfigMixin):
return_dict: bool = True,
train: bool = False,
) -> Union[FlaxUNet2DConditionOutput, Tuple]:
"""r
r"""
Args:
sample (`jnp.ndarray`): (channel, height, width) noisy inputs tensor
timestep (`jnp.ndarray` or `float` or `int`): timesteps
......
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