Fix: Remove _supported_args overwrite in AutoModeLForCausalLMWithValu…

…eHead
CarperAI · Jul 26, 2023 · eb95a42 · eb95a42
1 parent 5683470
commit eb95a42
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/trlx/models/modeling_ppo.py b/trlx/models/modeling_ppo.py
@@ -270,8 +270,7 @@ class AutoModelForCausalLMWithValueHead(PreTrainedModelWrapper):
 
     _auto_model_parent_class = transformers.AutoModelForCausalLM
     _supported_modules = ["v_head"]
-    _supported_args = ["peft_config"]
-    _supported_args = ["num_value_layers_unfrozen"]
+    _supported_args = ["peft_config", "num_value_layers_unfrozen"]
 
     def __init__(
         self,
@@ -329,6 +328,7 @@ def forward(
         else:
             outputs = self.base_model(**forward_kwargs)
 
+        # TODO: Apply PEFT to value branch
         if self.num_value_layers_unfrozen > 0:
             output_shape = outputs.hidden_states[-1].size()
             forward_kwargs.pop("input_ids", None)