from hf_olmo import OLMoForCausalLM, OLMoConfig
from transformers import AutoModel

AutoModel.register(OLMoConfig, OLMoForCausalLM)