Upload ConditionalUNet

config.json

@@ -2,6 +2,11 @@
   "architectures": [
     "ConditionalUNet"
   ],
+  "auto_map": {
+    "AutoConfig": "configuration_conditional_unet.ConditionalUNetConfig",
+    "AutoModelForImageClassification": "modeling_conditional_unet.ConditionalUNet"
+  },
+  "encoder_rep": "evanrsl/resnet-Alzheimer",
   "id2label": {
     "0": "LABEL_0",
     "1": "LABEL_1",
@@ -14,6 +19,7 @@
     "LABEL_2": 2,
     "LABEL_3": 3
   },
+  "model_type": "conditional-unet",
   "num_channels": 3,
   "torch_dtype": "float32",
   "transformers_version": "4.46.2"

model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:791ca306289c256b2cfe0260a5a6f2dad9b6f8c1348e93c7976734a4d327bca6
+oid sha256:77d484e051e9ea45cc8ebbd871d27d3f30a912bf39f4df7b46834f0c33135dbe
 size 293858844

modeling_conditional_unet.py

@@ -27,10 +27,11 @@ class UpSampleBlock(nn.Module):
         return x
 
 class ConditionalUNet(PreTrainedModel):
-    config_class = "ConditionalUNetConfig"
+    config_class = ConditionalUNetConfig
 
-    def __init__(self, config: ConditionalUNetConfig):
+    def __init__(self, config):
         super().__init__(config)
+        # self.config_class = config.__class__
         self.config = config
 
         self.encoder_rep = config.encoder_rep