Fine-tuning model does'nt learn targets

vision
1

I am trying to fine-tune a model from torch. hub but after training the loss goes down but the model doesn’t learn the target at all. Any hints on what the problem could be?

#Load pre-trained model and fine-tune

model = torch.hub.load('yangsenius/TransPose:main', 

                       'tpr_a4_256x192',

                       pretrained=True)

model.final_layer =  torch.nn.Sequential(torch.nn.Conv2d(256, 18,1))                                    

model = model.to(device)

#Get model summary

summary(model,

          input_size=(3, 256, 192),

          batch_size=1

          )

criterion = torch.nn.MSELoss(reduction="mean")

pretrain_part = [param for name, param in model.named_parameters()if 'final_layer' not in name]

# optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)

optimizer = torch.optim.Adam([ {'params': pretrain_part, 'lr':1e-5 },

                                                    {'params': model.final_layer.parameters(), 'lr': 1e-4}])