Question, which part do you implement flash-attention to accelerate?

[URRealHero]

You've mentioned flash-attn implementation, is there a chance to implement that into the diffusion training part? since I only found vae in this repo. Thanks a lot for your great work anyway

[URRealHero]

hmm, I think I found them. So, my current question is, which part of the data is used for training, and I have tried to decrease lr, but it seems the loss always goes back to 0.9 at a high-level. Is that a normal phenomenon?