Onnx model on GPU

[AlyaNovikova]

Hello! I used your code for the HAT-GAN model and further finetuned the model on my dataset. I converted the resulting pth model into onnx format. I tried to convert both with my own script and with the chaiNNer. But in both cases, when using the onnx model on the GPU, some layers are counted on the CPU and this causes a slowdown. I believe the problem is in the layers

self.residual_group = AttenBlocks(...)

It seems to me that they are counted on the CPU and not on the GPU, but I don’t understand why. Please tell me, maybe you know what the problem might be and how to rewrite the code so that it can all be used on the GPU. And maybe you or someone else has also encountered this problem