Upload model.py with huggingface_hub

model.py

@@ -4,7 +4,7 @@ import torch.nn as nn
 import torchvision
 import torch.nn.functional as F
 from torchvision.transforms import functional as FF
-from transformers import PretrainedConfig
+from transformers import PretrainedConfig, PreTrainedModel
 
 ######################################################################
 # Configuration File
@@ -223,10 +223,10 @@ class Decoder(nn.Module):
     x = self.double_conv(x)
     return x
 
-class VesselSegmentModel(nn.Module):
-  def __init__(self, config: VesselSegmentConfig):
-    super().__init__()
-    self.config = config
+class VesselSegmentModel(PreTrainedModel):
+  config_class = VesselSegmentConfig
+  def __init__(self, config: VesselSegmentConfig=VesselSegmentConfig()):
+    super().__init__(config)
     # image patch
     self.img_patch = ImagePatching(patch_size=config.patch_size)