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")