missing attribute

CarperAI · Aug 8, 2023 · 5cc4a17 · 5cc4a17
1 parent 9dae10c
commit 5cc4a17
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/trlx/models/modeling_ppo.py b/trlx/models/modeling_ppo.py
@@ -279,6 +279,7 @@ def __init__(
         num_value_layers_unfrozen=0,
     ):
         super().__init__(base_model, peft_config=peft_config)
+        self.num_value_layers_unfrozen = num_value_layers_unfrozen
         parameter = next(hf_get_lm_head(self.base_model).parameters())
         dtype = parameter.dtype
         device = parameter.device