Thanks for the code!
It looks like you are initializing some modules (e.g. nn.GroupNorm) inside your forward method.
This approach won’t push these modules onto the device, if you use model.cuda() or model.to(device).
Also, this will reset all parameters in each forward call (weight and bias), which is probably also not what you want.
Do you need and want to initialize them for each forward pass?