Fixing doc string

datamol-io · Aug 4, 2023 · da0d058 · da0d058
1 parent 6463e48
commit da0d058
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/graphium/finetuning/finetuning_architecture.py b/graphium/finetuning/finetuning_architecture.py
@@ -230,10 +230,10 @@ def overwrite_with_pretrained(
         Overwrite parameters shared between loaded and modified pretrained model
 
         Parameters:
-            pretrained_model:
+            pretrained_model: Model from GRAPHIUM_PRETRAINED_MODELS_DICT
             finetuning_module: Module to finetune from
             added_depth: Number of modified layers at the end of finetuning module
-            sub_module_from_pretrained: Optional submodule to finetune from FINETUNING_HEADS_DICT
+            sub_module_from_pretrained: Optional submodule to finetune from
         """
         module_map = self.net._module_map
         module_map_from_pretrained = pretrained_model._module_map
@@ -292,7 +292,7 @@ def __init__(self, finetuning_head_kwargs: Dict[str, Any]):
 
         Parameters:
 
-            finetuning_head_kwargs: Key-word arguments needed to instantiate a custom (or existing) finetuning head from
+            finetuning_head_kwargs: Key-word arguments needed to instantiate a custom (or existing) finetuning head from FINETUNING_HEADS_DICT
 
         """