🧩 Add optimizer_cls_and_kwargs attribute to PPOTrainer and `RLOOT…

…rainer` (#2302)
huggingface · Oct 31, 2024 · d57a181 · d57a181
1 parent 73c3970
commit d57a181
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 0 deletions.
diff --git a/trl/trainer/ppo_trainer.py b/trl/trainer/ppo_trainer.py
@@ -131,6 +131,7 @@ def __init__(
         self.data_collator = data_collator
         self.eval_dataset = eval_dataset
         self.optimizer, self.lr_scheduler = optimizers
+        self.optimizer_cls_and_kwargs = None # needed for transformers >= 4.47
 
         #########
         # calculate various batch sizes

diff --git a/trl/trainer/rloo_trainer.py b/trl/trainer/rloo_trainer.py
@@ -110,6 +110,7 @@ def __init__(
         self.data_collator = data_collator
         self.eval_dataset = eval_dataset
         self.optimizer, self.lr_scheduler = optimizers
+        self.optimizer_cls_and_kwargs = None  # needed for transformers >= 4.47
 
         #########
         # calculate various batch sizes