Loss not decreasing with QAT applied

Hello everyone, I wanted to ask for help regarding to QAT. I am using an open source repository for pose estimation, CenterNet GitHub - xingyizhou/CenterNet: Object detection, 3D detection, and pose estimation using center point detection:.
I am using essentially all their structure except that I created my own ResNet18 architecture and trained a model with that architecture. The model without QAT is working but when I try to add the QAT my loss is not decreasing and I am not sure where is the correct place to convert it.
At the beginning I thought that the “convert” step was my problem, because when I trained the whole model my accuracy dropped to 0 and the loss was oscillating in the same values. but after I created an unit test for my loss calculation with QAT applied the loss and shapes are correct. Essentially I am doing the same in my unit test as in the main, except that instead of using the original dataset I generated a dataset with random numbers(in my unit test). In both options I have my QAT configuration applied without the “convert” step. Meaning that there is something in the middle that is making my unit test work correctly and CenterNet model not, but I cannot find it. You will notice that the QAT is being apply just in my fist layers before deconvolution.

Could you guys help me trying to figured out? I will add in a private colab(unfortunately I cannot publish it) my modified main with the QAT configuration and my architecture. For the moment I just added @jerryzh168 as collaborator, but if someone else could help me too, I will be grateful and I can add them too.Google Colab

Thank you very much in advance