Training notebook does not seem to work with 0.3.0

[patrickvonplaten]

Describe the bug

The training notebook doesn't seem to work with 0.3.0

Reproduction

Running this google colab: https://colab.research.google.com/drive/1wAOAo9LCjWukVp6QBwpw6EwoVhsdyuNj?usp=sharing

Logs

AttributeError                            Traceback (most recent call last)
<ipython-input-51-4f9138074009> in <module>
      2 args = (config, model, noise_scheduler, optimizer, train_dataloader, lr_scheduler)
      3 
----> 4 notebook_launcher(train_loop, args, num_processes=1)

5 frames
/usr/local/lib/python3.7/dist-packages/diffusers/schedulers/scheduling_ddpm.py in step(self, model_output, timestep, sample, predict_epsilon, generator, return_dict)
    202         t = timestep
    203 
--> 204         if model_output.shape[1] == sample.shape[1] * 2 and self.variance_type in ["learned", "learned_range"]:
    205             model_output, predicted_variance = torch.split(model_output, sample.shape[1], dim=1)
    206         else:

AttributeError: 'dict' object has no attribute 'shape'

System Info

!pip install diffusers[training]==0.3.0

https://colab.research.google.com/drive/1wAOAo9LCjWukVp6QBwpw6EwoVhsdyuNj#scrollTo=r5PM6vOQPISl

[patrickvonplaten]

@anton-l could you take a look here?

[anton-l]

Updated the notebook in huggingface/notebooks@c0cff84
A more flexible/permanent fix is in #570