_dict: bool = True

)
→
UNet2DOutput or tuple

Parameters 

sample (torch.FloatTensor) — (batch, channel, height, width) 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 UNet2DOutput instead of a plain tuple.


Returns

UNet2DOutput or tuple



UNet2DOutput if return_dict is True,
otherwise a tuple. When returning a tuple, the fi