Updated code

model_dir/config.json

@@ -1,23 +0,0 @@
-{
-  "_name_or_path": "shravvvv/SAG-ViT",
-  "architectures": [
-    "SAGViTClassifier"
-  ],
-  "d_model": 64,
-  "dim_feedforward": 64,
-  "gcn_hidden": 128,
-  "gcn_out": 64,
-  "hidden_mlp_features": 64,
-  "in_channels": 2560,
-  "model_type": "sagvit",
-  "nhead": 4,
-  "num_classes": 10,
-  "num_layers": 2,
-  "patch_size": [
-    4,
-    4
-  ],
-  "torch_dtype": "float32",
-  "transformers_version": "4.47.0",
-  "use_safetensors": true
-}

model_dir/pytorch_model.bin

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:3472578af5e4fd2f9644e94f8502895fc007abea1e6880364c0480b588c474a0
-size 32491922

push_model_to_hfhub.py

@@ -1,10 +1,13 @@
 from transformers import AutoConfig, AutoModel
 from modeling_sagvit import SAGViTClassifier, SAGViTConfig
 
+
+AutoConfig.register("sagvit", SAGViTConfig)
+AutoModel.register(SAGViTConfig, SAGViTClassifier)
+
 # Load config and model
-config = AutoConfig.from_pretrained("shravvvv/SAG-ViT")
+config = SAGViTConfig()
 model = SAGViTClassifier(config)
 
-# Save model locally before pushing
-model.save_pretrained("./model_dir", safe_serialization=False)  # Save in PyTorch format
+# Push model and code
 model.push_to_hub("shravvvv/SAG-ViT")

register_model.py

@@ -1,14 +1,8 @@
 from transformers import AutoConfig, AutoModel
-from transformers.models.auto.configuration_auto import CONFIG_MAPPING
-from transformers.models.auto.modeling_auto import MODEL_MAPPING
+from modeling_sagvit import SAGViTConfig, SAGViTClassifier
 
-from sagvit_config import SAGViTConfig
-from modeling_sagvit import SAGViTClassifier
-
-# Register the configuration
-CONFIG_MAPPING.register("sagvit", SAGViTConfig)
-
-# Register the model
-MODEL_MAPPING.register(SAGViTConfig, SAGViTClassifier)
-
-print("Registered model successfully...")
+# Register Custom Model and Config
+print("Registering model")
+AutoConfig.register("sagvit", SAGViTConfig)
+AutoModel.register(SAGViTConfig, SAGViTClassifier)
+print("Registration complete")