Update modeling_auristream.py

modeling_auristream.py

@@ -12,9 +12,11 @@ import numpy as np
 from huggingface_hub import PyTorchModelHubMixin
 from transformers.modeling_outputs import BaseModelOutput, CausalLMOutput
 from transformers import PreTrainedModel
+from .configuration_auristream import AuriStreamConfig
 
 
 class AuriStream(PreTrainedModel):
+    config_class = AuriStreamConfig
 
     def __init__(self, config):
         super().__init__(config)