add model

config.json

@@ -4,7 +4,7 @@
   ],
   "auto_map": {
     "AutoConfig": "config.CNNConfig",
-    "AutoModel": "__main__.CNNModel"
+    "AutoModel": "model.CNNModel"
   },
   "num_classes": 10,
   "torch_dtype": "float32",

model.py

@@ -1,4 +1,4 @@
-from config import CNNConfig
+from .config import CNNConfig
 from transformers import PreTrainedModel
 
 
@@ -71,20 +71,5 @@ class CNNModel(PreTrainedModel):
     def forward(self, tensor):
         return self.model(tensor)
 
-if __name__ == "__main__":
-    weights = "k_cifar10/2wye5gyi/checkpoints/epoch=599-step=187800.ckpt"
-
-    CNNConfig.register_for_auto_class()
-    CNNModel.register_for_auto_class("AutoModel")
-
-    pretrained = CNN.load_from_checkpoint(weights)
-
-    config = CNNConfig.from_pretrained("custom_cnn_10")
-    my_cnn = CNNModel(config)
-    my_cnn.model.load_state_dict(pretrained.state_dict())
-
-    my_cnn.push_to_hub("k_cnn_cifar10")
-
-