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