model init when args are given (huggingface#1413)

Co-authored-by: Lewis Tunstall <lewis.c.tunstall@gmail.com>
fe1ixxu · Mar 15, 2024 · cd0b4c9 · cd0b4c9
1 parent 291adae
commit cd0b4c9
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/trl/trainer/kto_trainer.py b/trl/trainer/kto_trainer.py
@@ -126,13 +126,17 @@ def __init__(
             model_init_kwargs = {}
         elif not isinstance(model, str):
             raise ValueError("You passed model_kwargs to the KTOTrainer. But your model is already instantiated.")
+        else:
+            model_init_kwargs = args.model_init_kwargs
 
         if args.ref_model_init_kwargs is None:
             ref_model_init_kwargs = {}
         elif not isinstance(ref_model, str):
             raise ValueError(
                 "You passed ref_model_kwargs to the KTOTrainer. But your ref_model is already instantiated."
             )
+        else:
+            ref_model_init_kwargs = args.ref_model_init_kwargs
 
         if isinstance(model, str):
             warnings.warn(