ape matches the model output shape:


	Copied
>>> sample_image = dataset[0]["images"].unsqueeze(0)
>>> print("Input shape:", sample_image.shape)
Input shape: torch.Size([1, 3, 128, 128])

>>> print("Output shape:", model(sample_image, timestep=0).sample.shape)
Output shape: torch.Size([1, 3, 128, 128])
Great! Next, you’ll need a scheduler to add some noise to the image.

Create a scheduler
	
The scheduler behaves differently depending on whethe