Update model.py

model.py

@@ -114,4 +114,10 @@ class CausalLMForRegression(nn.Module):
         prepare_inputs_for_generation, etc.) keep working. We simply forward all
         arguments to the underlying causal‑LM that already implements them.
         """
-        return self.model.generate(*args, **kwargs)
+        return self.model.generate(*args, **kwargs)
+
+    def prepare_inputs_for_generation(self, *args, **kwargs):
+        """
+        Same here: to be able to load the model with AutoModelForCausalLM, we have to forward this method
+        """
+        return self.model.prepare_inputs_for_generation(*args, **kwargs)