| from transformers import AutoConfig, AutoModel | |
| from modeling_sagvit import SAGViTConfig, SAGViTClassifier | |
| # Register Custom Model and Config | |
| print("Registering model") | |
| AutoConfig.register("sagvit", SAGViTConfig) | |
| AutoModel.register(SAGViTConfig, SAGViTClassifier) | |
| print("Registration complete") | |