Question on training Unet2DModel to predict the clean instead of noise

Thanks for the notebooks.

I have one comment
In this file https://github.com/huggingface/diffusion-models-class/blob/main/unit1/02_diffusion_models_from_scratch.ipynb
```python
# The training loop
for epoch in range(n_epochs):

for x, y in train_dataloader:

# Get some data and prepare the corrupted version
x = x.to(device) # Data on the GPU
noise_amount = torch.rand(x.shape[0]).to(device) # Pick random noise amounts
noisy_x = corrupt(x, noise_amount) # Create our noisy x

# Get the model prediction
pred = net(noisy_x, 0).sample #<<< Using timestep 0 always, adding .sample

# Calculate the loss
loss = loss_fn(pred, x) # How close is the output to the true 'clean' x?

# Backprop and update the params:
opt.zero_grad()
loss.backward()
opt.step()

# Store the loss for later
losses.append(loss.item())

# Print our the average of the loss values for this epoch:
avg_loss = sum(losses[-len(train_dataloader):])/len(train_dataloader)
print(f'Finished epoch {epoch}. Average loss for this epoch: {avg_loss:05f}')
```
In case I want to make the network predict the clean images, should the pred formula be changed to this by attaching noise_amount?
```python
# Get the model prediction
pred = net(noisy_x, noise_amount).sample #<<< Using timestep 0 always, adding .sample
```