Fix: Missing num_value_layers_unfrozen init arg

CarperAI · Jul 26, 2023 · 5683470 · 5683470
1 parent a875c4a
commit 5683470
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/trlx/trainer/accelerate_ppo_trainer.py b/trlx/trainer/accelerate_ppo_trainer.py
@@ -115,6 +115,7 @@ def get_arch(self, config: TRLConfig):
         return from_fn(
             config.model.model_path,
             num_layers_unfrozen=config.model.num_layers_unfrozen,
+            num_value_layers_unfrozen=config.method.num_value_layers_unfrozen,
             peft_config=self.config.model.peft_config,
         )