Update model.py

model.py

@@ -4,9 +4,14 @@ from transformers import AutoModelForCausalLM, AutoTokenizer
 import torch
 import torch.nn as nn
 
+from transformers.models.qwen3 import Qwen3Config
+
 
 # Define a custom model that wraps a causal LM and adds a regression head
 class CausalLMForRegression(nn.Module):
+    config_class = Qwen3Config
+    base_model_prefix = "model"
+    
     def __init__(self, model_name):
         super().__init__()
         # Load the causal LM with hidden states enabled