Upload DiskForKeypointDetection

modeling_disk.py

@@ -2,10 +2,10 @@ import kornia
 import torch
 
 from .configuration_disk import DiskConfig
-from transformers import AutoConfig, AutoModelForKeypointDetection, PreTrainedModel
 from transformers.models.superpoint.modeling_superpoint import (
     SuperPointKeypointDescriptionOutput,
 )
+from transformers import PreTrainedModel
 
 
 class DiskForKeypointDetection(PreTrainedModel):
@@ -15,7 +15,7 @@ class DiskForKeypointDetection(PreTrainedModel):
         super().__init__(config)
 
         self.config = config
-        self.model = kornia.feature.DISK.from_pretrained(self.config.weights)
+        self.model = kornia.feature.DISK(self.config.descriptor_decoder_dim)
 
     def forward(
         self, pixel_values: torch.Tensor