Code fails the backward pass

[RishikeshMagar]

Thanks for creating this great repo.

I was trying to run your code in our training setup, however the code failed in the backward pass. The forward pass works fine, however in the backward pass the code fails at this at this line -----> assert do.is_contiguous()

flashattention2-custom-mask/fa2_custom_mask/fa2_custom_mask.py

Line 61 in 8170916

assert do.is_contiguous()