reorder model wrap + bug fix

[prathikr]

Fixes return_dict bug and resolves the following pickling error when using UNet2DModel in instantiation of ema_model:

Traceback (most recent call last):
  File "train_unconditional_ort.py", line 452, in <module>
    main(args)
  File "train_unconditional_ort.py", line 327, in main
    ema_model = EMAModel(
  File "/home/prathikrao/diffusers/src/diffusers/training_utils.py", line 69, in __init__
    self.averaged_model = copy.deepcopy(model).eval()
  File "/opt/conda/envs/ptca/lib/python3.8/copy.py", line 172, in deepcopy
    y = _reconstruct(x, memo, *rv)
  File "/opt/conda/envs/ptca/lib/python3.8/copy.py", line 270, in _reconstruct
    state = deepcopy(state, memo)
  File "/opt/conda/envs/ptca/lib/python3.8/copy.py", line 146, in deepcopy
    y = copier(x, memo)
  File "/opt/conda/envs/ptca/lib/python3.8/copy.py", line 230, in _deepcopy_dict
    y[deepcopy(key, memo)] = deepcopy(value, memo)
  File "/opt/conda/envs/ptca/lib/python3.8/copy.py", line 161, in deepcopy
    rv = reductor(4)
  File "/opt/conda/envs/ptca/lib/python3.8/site-packages/accelerate/utils/operations.py", line 493, in __getstate__
    raise pickle.PicklingError(
_pickle.PicklingError: Cannot pickle a prepared model with automatic mixed precision, please unwrap the model with `Accelerator.unwrap_model(model)` before pickling it.

[prathikr]

@anton-l could I please get a review on this fix?

[HuggingFaceDocBuilderDev]

The documentation is not available anymore as the PR was closed or merged.

[patil-suraj]

LGTM, thanks a lot for the fix!