Update model.py

model.py

@@ -88,26 +88,33 @@ class CausalLMForRegression(nn.Module):
 
     @classmethod
     def from_pretrained(cls, pretrained_model_name_or_path, *model_args, **kwargs):
-        # make sure hidden states are returned
-        kwargs.setdefault("output_hidden_states", True)
+        config = AutoConfig.from_pretrained(pretrained_model_name_or_path, **kwargs)
+        config.output_hidden_states = True
 
         base_model = Qwen3ForCausalLM.from_pretrained(
-            pretrained_model_name_or_path, *model_args, **kwargs
+            pretrained_model_name_or_path,
+            *model_args,
+            config=config,            
+            **{k: v for k, v in kwargs.items()
+               if k not in config.to_dict()}  
         )
 
         instance = cls.__new__(cls)
         nn.Module.__init__(instance)
+
         instance.model = base_model
-        instance.regression_head = nn.Linear(base_model.config.hidden_size, 1)
+        instance.regression_head = nn.Linear(config.hidden_size, 1)
+        instance._keys_to_ignore_on_save = []
 
-        head_path = os.path.join(pretrained_model_name_or_path, "regression_head.bin")
+        head_path = os.path.join(pretrained_model_name_or_path,
+                                 "regression_head.bin")
         if os.path.exists(head_path):
             instance.regression_head.load_state_dict(
                 torch.load(head_path, map_location="cpu")
             )
         else:
             print("No regression head found – initialising randomly.")
-        instance._keys_to_ignore_on_save = []
+
         return instance
 
     @torch.no_grad()