Try and get this custom stuff working

modeling_diva.py

@@ -11,11 +11,11 @@ import torch.nn.functional as F
 from datasets import Audio
 from safetensors.torch import load, load_model
 from torch import nn
+from .configuring_diva import DiVAConfig
 from transformers import (
     AutoProcessor,
     AutoTokenizer,
     LlamaForCausalLM,
-    PretrainedConfig,
     PreTrainedModel,
     WhisperForConditionalGeneration,
 )
@@ -44,10 +44,12 @@ class WhisperConnector(nn.Module):
 
 
 class DiVAModel(PreTrainedModel):
+    config_class = DiVAConfig
+
     def __init__(
         self, via_path=None, config_dict={}, device_map=None, speech_encoder_device=None
     ):
-        super().__init__(PretrainedConfig.from_dict(config_dict))
+        super().__init__(DiVAConfig.from_dict(config_dict))
         if speech_encoder_device is None:
             speech_encoder_device = "cuda:0"
         whisper = WhisperForConditionalGeneration.from_pretrained(