doc string args shape fix (#1243)
* doc string args shape fix * fix styling
This commit is contained in:
parent
f1fcfdeec5
commit
aa5c4c2609
|
@ -251,7 +251,8 @@ class UNet2DConditionModel(ModelMixin, ConfigMixin):
|
||||||
Args:
|
Args:
|
||||||
sample (`torch.FloatTensor`): (batch, channel, height, width) noisy inputs tensor
|
sample (`torch.FloatTensor`): (batch, channel, height, width) noisy inputs tensor
|
||||||
timestep (`torch.FloatTensor` or `float` or `int`): (batch) timesteps
|
timestep (`torch.FloatTensor` or `float` or `int`): (batch) timesteps
|
||||||
encoder_hidden_states (`torch.FloatTensor`): (batch, channel, height, width) encoder hidden states
|
encoder_hidden_states (`torch.FloatTensor`):
|
||||||
|
(batch_size, sequence_length, hidden_size) encoder hidden states
|
||||||
return_dict (`bool`, *optional*, defaults to `True`):
|
return_dict (`bool`, *optional*, defaults to `True`):
|
||||||
Whether or not to return a [`models.unet_2d_condition.UNet2DConditionOutput`] instead of a plain tuple.
|
Whether or not to return a [`models.unet_2d_condition.UNet2DConditionOutput`] instead of a plain tuple.
|
||||||
|
|
||||||
|
|
|
@ -230,9 +230,9 @@ class FlaxUNet2DConditionModel(nn.Module, FlaxModelMixin, ConfigMixin):
|
||||||
) -> Union[FlaxUNet2DConditionOutput, Tuple]:
|
) -> Union[FlaxUNet2DConditionOutput, Tuple]:
|
||||||
r"""
|
r"""
|
||||||
Args:
|
Args:
|
||||||
sample (`jnp.ndarray`): (channel, height, width) noisy inputs tensor
|
sample (`jnp.ndarray`): (batch, channel, height, width) noisy inputs tensor
|
||||||
timestep (`jnp.ndarray` or `float` or `int`): timesteps
|
timestep (`jnp.ndarray` or `float` or `int`): timesteps
|
||||||
encoder_hidden_states (`jnp.ndarray`): (channel, height, width) encoder hidden states
|
encoder_hidden_states (`jnp.ndarray`): (batch_size, sequence_length, hidden_size) encoder hidden states
|
||||||
return_dict (`bool`, *optional*, defaults to `True`):
|
return_dict (`bool`, *optional*, defaults to `True`):
|
||||||
Whether or not to return a [`models.unet_2d_condition_flax.FlaxUNet2DConditionOutput`] instead of a
|
Whether or not to return a [`models.unet_2d_condition_flax.FlaxUNet2DConditionOutput`] instead of a
|
||||||
plain tuple.
|
plain tuple.
|
||||||
|
|
Loading…
Reference in New Issue