Repository: huggingface/diffusers

Issue:
Example unconditional training script does not work
### Describe the bug

When running the [example unconditional training script](https://github.com/huggingface/diffusers/blob/main/examples/unconditional_image_generation/train_unconditional.py)
on GPU using `--mixed_precision=fp16`, the autocasting of accelerate will break:

```
Traceback (most recent call last):
  File "<path>\train_unconditional.py", line 251, in <module>
    main(args)
  File "<path>\train_unconditional.py", line 145, in main
    loss = F.mse_loss(noise_pred, noise)
  File "<lib-path>\lib\site-packages\torch\nn\functional.py", line 3269, in mse_loss
    if not (target.size() == input.size()):
AttributeError: 'dict' object has no attribute 'size'
```

This happens due to the fact that in `accelerate.utils.operations.recursively_apply` the "Mapping" class is handled using a dict initialization (=> https://github.com/huggingface/accelerate/blob/main/src/accelerate/utils/operations.py#L87). This breaks the UNet2DOutput class, because it will wrap the output in a dict (because UNet2DOutput does inherit from OrderedDict).
i.e. `UNet2DOutput(sample=sample)` will become equivalent to `UNet2DOutput(sample={'sample': sample})`.

now this could be changed in the train_unconditional.py. unfortunately this just moves the problem, as it will fail in a later stage in library code that isn't accessible to the user.

### Reproduction

Run the example script with mixed-precision set to fp16:
```
train_unconditional.py --dataset_name="huggan/pokemon" --output_dir="ddpm-test-64" --mixed_precision=fp16
```

### Logs

_No response_

### System Info

- `diffusers` version: 0.3.0
- Platform: Windows-10-10.0.22000-SP0
- Python version: 3.9.13
- PyTorch version (GPU?): 1.12.1+cu116 (True)
- Huggingface_hub version: 0.9.1
- Transformers version: 4.21.3
- Using GPU in script?: yes
- Using distributed or parallel set-up in script?: no

some comments:
*  i'm running on windows 11, wrongly reported as 10 using `diffusers-cli env`
* `accelerate` version: 0.12.0


Discussion:
Thanks for reporting @c-goldschmidt! Looking into it as part of https://github.com/huggingface/diffusers/issues/434/
I am facing the same error. After installing the latest diffusers from GitHub

```
pip install git+https://github.com/huggingface/diffusers.git
```

The original issue `AttributeError: 'dict' object has no attribute 'size'` is gone, but another issue appears

```
Traceback (most recent call last):
  File "train_unconditional.py", line 249, in <module>
    main(args)
  File "train_unconditional.py", line 146, in main
    accelerator.clip_grad_norm_(model.parameters(), 1.0)
  File ".../miniconda3/envs/diffusion/lib/python3.8/site-packages/accelerate/accelerator.py", line 920, in clip_grad_norm_
    self.unscale_gradients()
  File ".../miniconda3/envs/diffusion/lib/python3.8/site-packages/accelerate/accelerator.py", line 904, in unscale_gradients
    self.scaler.unscale_(opt)
  File ".../miniconda3/envs/diffusion/lib/python3.8/site-packages/torch/cuda/amp/grad_scaler.py", line 270, in unscale_
    raise RuntimeError("unscale_() has already been called on this optimizer since the last update().")
RuntimeError: unscale_() has already been called on this optimizer since the last update().
```
After applying #686, the problem is resolved.
Thanks for reporting @gau-nernst ! Would you like to open a PR to apply this fix to the `train_unconditional.py` script ?
Also cc @anton-l for `train_unconditional.py`
> Thanks for reporting @gau-nernst ! Would you like to open a PR to apply this fix to the `train_unconditional.py` script ?

After applying your fix at #686 (before it was merged), the problem is resolved. So I think it should be good now.

Resolving patch:
diff --git a/examples/unconditional_image_generation/train_unconditional.py b/examples/unconditional_image_generation/train_unconditional.py
--- a/examples/unconditional_image_generation/train_unconditional.py
+++ b/examples/unconditional_image_generation/train_unconditional.py
@@ -143,7 +143,8 @@ def transforms(examples):
                 loss = F.mse_loss(noise_pred, noise)
                 accelerator.backward(loss)
 
-                accelerator.clip_grad_norm_(model.parameters(), 1.0)
+                if accelerator.sync_gradients:
+                    accelerator.clip_grad_norm_(model.parameters(), 1.0)
                 optimizer.step()
                 lr_scheduler.step()
                 if args.use_ema:


Test patch: