Not work for Swin_T model

Thank you guys for your library. It's a quite convenient tool.

However, when I want to implement simplify with Swin_T model, it seems it does not work.

To illustrate the problem, I've shared a simplified version of the code in the following link. Could you kindly review it and assist me in identifying a solution to rectify this issue? Your support would be greatly appreciated.
https://partage.imt.fr/index.php/s/9jemn7WfkBWS5tx

Thanks a lot.