Upload model

modeling_rf_detr.py

@@ -2,7 +2,7 @@ from dataclasses import dataclass
 from typing import List, Dict
 
 import torch
-from torchvision.transforms import Resize
+from torchvision.transforms import Resize, InterpolationMode
 from transformers import PreTrainedModel
 from transformers.utils import ModelOutput
 from rfdetr import RFDETRBase, RFDETRLarge
@@ -118,7 +118,7 @@ class RFDetrModelForObjectDetection(PreTrainedModel):
                 label["labels"] = label["labels"].to(self.config.device)
 
     def forward(self, pixel_values: torch.Tensor, pixel_mask: torch.Tensor, labels=None, **kwargs) -> ModelOutput:
-        resize = Resize((self.config.resolution, self.config.resolution), antialias=False) # antialias set to False for onnx export
+        resize = Resize((self.config.resolution, self.config.resolution), interpolation=InterpolationMode.NEAREST) # interpolation mode set to nearest for onnx export
 
         if labels is not None:
             self.validate_labels(labels)