__init__.py

from .configuration_custom_gpt import CustomGPTConfig
from .custom_gpt import CustomGPT

CustomGPTConfig.register_for_auto_class()
CustomGPT.register_for_auto_class("AutoModelForCausalLM")