__init__.py

# __init__.py

from transformers import AutoConfig, AutoModel
from .custom_model2 import CustomConfig, CustomModel

# Register the custom classes
AutoConfig.register("custom_model", CustomConfig)
AutoModel.register(CustomConfig, CustomModel)