Use chunkwise kernel with VisionLSTM?

[NaderEK]

Dears,

I hope you’re all doing well! First off, I want to say a big thank you for the amazing work you’re doing.

I’m currently diving into a research project where I'm using VisionLSTM for image restoration, but I’ve hit a bit of a snag. I'm trying to work with high-resolution images, but I'm running into memory issues since I’m not dividing them into smaller patches. Ideally, I want to restore them in their original sizes without much down-sampling and up-sampling.

Since the parallel kernel scales quadratically with sequence length, do you think switching to a "chunkwise" backend instead of the "parallel" backend option would help solve this problem?

Thanks so much for your help!