File size: 239 Bytes
2fe4bd1 | 1 2 3 4 5 6 | from transformers import AutoConfig, AutoModel
from src.models.neuroblast import NeuroBLASTConfig, NeuroBLASTForCausalLM
AutoConfig.register("neuroblast", NeuroBLASTForCausalLM)
AutoModel.register(NeuroBLASTConfig, NeuroBLASTForCausalLM)
|