Update modeling.py

modeling.py

@@ -1,12 +1,14 @@
+
 import torch
 import torch.nn as nn
 import timm
 from huggingface_hub import PyTorchModelHubMixin
 
 class KeypointModel(nn.Module, PyTorchModelHubMixin):
-    def __init__(self, heatmap_size=(512, 1024), **kwargs):
+    def __init__(self, config, **kwargs):
         super().__init__()
-        config_heatmap_size = kwargs.get("config", {}).get("heatmap_size", heatmap_size)
+
+        upsample_size = config.heatmap_size
 
         backbone = timm.create_model('convnextv2_base.fcmae_ft_in22k_in1k_384', pretrained=False)
 
@@ -15,7 +17,7 @@ class KeypointModel(nn.Module, PyTorchModelHubMixin):
         self.head = nn.Sequential(
             nn.Conv2d(in_channels, 256, kernel_size=3, padding=1),
             nn.ReLU(inplace=True),
-            nn.Upsample(size=config_heatmap_size, mode='bilinear', align_corners=False),
+            nn.Upsample(size=upsample_size, mode='bilinear', align_corners=False),
             nn.Conv2d(256, 1, kernel_size=1)
         )