thank you for your quick response .
I think the issue is coming when I train the model somehow backbone is getting updated even though it is in eval mode.
class Net(nn.Module):

  def __init__(self):
      self.backbone = timm.create_model('resnet50', num_classes=0).eval()
      self.head = nn.Linear(2048, num_classes)

  def forward(self, x):
      with torch.no_grad():
         x = self.backbone(x)
      x = self.head(x)
   return x

def train(net, dataloader):
    """model trained with crossentropy here"""
  

After training, what I want is to only save the head. I will do this with a couple of tasks and use a single pretrained frozen backbone for extracting features and use the head of each task for getting the probabilities.