SPLIT_EINSUM - Kernel Panic when testing UNET created with height 96 and width 64 and SPLIT_EINSUM

[littleowl]

I've experienced multiple kernel panics trying to create a model with latent space of 64w-96h. I was able to do this without much problems when --attention-implementation ORIGINAL was used and was able to output the correct image. However, when SPLIT_EINSUM is used I'll get kernel panics. This will happen when setting the flag in python --check-output-correctness but it will also happen in the swift CLI.

Again, I was able to get this to work with ORIGINAL and not setting .all but setting .cpuAndGPU. I'm on a M1 Max with 64gigs memory. If set to .all with the ORIGINAL attention, the program will crash.