this is my forward definition -
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