Replace SE with SA layer?

[MaxPolak97]

Thank you for this great repository an ShuffleNetv2+ is really awesome 😃. I wonder whether you can just easily replace the Squeeze-and-Excitation (SE) layer with the Shuffle Attention (SA) layer as proposed in https://blog.paperspace.com/shuffle-attention-sanet/ which actually uses a trick of ShuffleNetv2 to combine both channel and spatial attention to efficiently improve performance. Is this possible or are there other things I should consider?

The code of the SA layer can be found in the hyperlink.

Thanks in advance!