RuntimeError: CUDA illegal memory access

[karnikram]

Hi,

Thank you for your work and code, this is really cool!

I've been running some experiments with your code and I noticed the gradient_bias option inside the Patchifier network to bias patch selection towards regions with high gradient.

However when I enable this option the training exits with an error:

RuntimeError: CUDA error: an illegal memory access was encountered

I think this is happening because of a shape mismatch inside the patchify_forward_kernel cuda kernel coming from the gradient images, but I wanted to check if you are aware of this error?

[lahavlipson]

Thanks for pointing this out, I've just added a fix.

The gradient_bias wasn't tested during training, and during inference it doesn't seem to perform much better than randomly sampling patch centroids.