fix peft_config type (#883)

Co-authored-by: wanglei.w <[email protected]>
huggingface · Oct 18, 2023 · 122edc8 · 122edc8
1 parent f91fb2b
commit 122edc8
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/trl/trainer/sft_trainer.py b/trl/trainer/sft_trainer.py
@@ -114,7 +114,7 @@ def __init__(
         callbacks: Optional[List[TrainerCallback]] = None,
         optimizers: Tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LambdaLR] = (None, None),
         preprocess_logits_for_metrics: Optional[Callable[[torch.Tensor, torch.Tensor], torch.Tensor]] = None,
-        peft_config: Optional[Dict] = None,
+        peft_config: Optional[PeftConfig] = None,
         dataset_text_field: Optional[str] = None,
         packing: Optional[bool] = False,
         formatting_func: Optional[Callable] = None,