type hints for load peft

mosaicml · Aug 11, 2023 · 47d9811 · 47d9811
1 parent ade0f2a
commit 47d9811
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/scripts/eval/eval.py b/scripts/eval/eval.py
@@ -26,7 +26,8 @@
 from llmfoundry.utils.config_utils import process_init_device
 
 
-def load_peft_model(model_cfg, tokenizer, num_retries):
+def load_peft_model(model_cfg: DictConfig, tokenizer: PreTrainedTokenizerBase,
+                    num_retries: int) -> Optional[ComposerModel]:
     try:
         from peft import PeftModel
     except ImportError as e: