@ptrblck i need 1 help from you!
i tried your bn to gn swapping code and it works but here is what happening with me :

without groupnormalization swapping my model works fine
when i swap bn layers with gn(like you did in above message) i see qwk score getting 0.0 sometimes but with batchnorm layers i never get qwk score,why is this happening?

let me repeat again : “everything works fine when i don’t swap bn layers with gn layers”
when i use gn layers my model starts performing very badly but groupnorm paper is telling different story!! what could be the reason for this? any suggestion or groupnorm modification for solving this issue?