Problem in backward

[haoweiz23]

Describe the bug

Hi, I modified the U-Net forward function to use the "wrapped_forward" function. However, I encountered an issue when attempting to backpropagate the loss. I adapted the training code from "examples/train_dreambooth.py". The error log is provided below. It appears to be related to gradient checkpointing, but I couldn't find any useful information to resolve it.

Reproduction

Logs

No response

System Info

Copy-and-paste the text below in your GitHub issue and FILL OUT the two last points.

• diffusers version: 0.28.0.dev0
• Platform: Linux-5.4.0-148-generic-x86_64-with-glibc2.31
• Python version: 3.10.6
• PyTorch version (GPU?): 1.13.1+cu117 (True)
• Huggingface_hub version: 0.20.2
• Transformers version: 4.39.0
• Accelerate version: 0.25.0
• xFormers version: not installed
• Using GPU in script?:
• Using distributed or parallel set-up in script?:

Who can help?

No response

[sayakpaul]

The error seems to be coming from a piece code that lies outside of the library. So, I suggest you open a discussion for this instead of an issue.

[tolgacangoz]

Hi @haoweiz23,
To be honest, I couldn't understand what is going on just to see these two pictures. Could you please share your modified examples/train_dreambooth.py and U-Net files and give details about the intention of your wrapping?