Retraining head for ssdlite-mobilenetv3

Hi! I have a pretrained ssdlite320_mobilenet_v3_large model with a classification/regression head.
The model performs object detection + classification among several tasks. However, now I want to finetune the model so that it works on a reduced number of classes.

How should I go about retraining/finetuning the head part only? Is retraining even required?

I am trying something like this, but the output dimensions during training seem to be off.

 model_old = SSDMobilenet(num_classes=20) # 
 model_old.load_state_dict(model_path, map_location=self.device)
  # New model with user-defined number of classes
  model = SSDMobilenet(num_classes=4) 
  # Use old model backbone with new classifier
  model.backbone = model_old.torchvision_model.backbone
  for param in model.parameters():
      param.requires_grad = False
  for param in model.torchvision_model.head.classification_head.parameters():
      param.requires_grad = True