Upload model.py with huggingface_hub

model.py

@@ -4,12 +4,14 @@ import torch.nn as nn
 import torchvision
 import torch.nn.functional as F
 from torchvision.transforms import functional as FF
+from transformers import PretrainedConfig
 
 ######################################################################
 # Configuration File
 ######################################################################
-class VesselSegmentConfig:
-  def __init__(self, num_classes, input_channels, image_size, features, attention_dims, patch_size, batch_size):
+class VesselSegmentConfig(PretrainedConfig):
+  model_type = "SegformerForSemanticSegmentation"
+  def __init__(self, num_classes, input_channels, image_size, features, attention_dims, patch_size, batch_size, **kwargs):
     self.num_classes = num_classes
     self.input_channels = input_channels
     self.image_size = image_size
@@ -17,6 +19,7 @@ class VesselSegmentConfig:
     self.attention_dims = attention_dims
     self.patch_size = patch_size
     self.batch_size = batch_size
+    super().__init__(**kwargs)
 
 ######################################################################
 # IMAGE DOWN SAMPLING