add missing params to docstring

basf · Apr 22, 2024 · 0784326 · 0784326
1 parent 0cf0eda
commit 0784326
Showing 1 changed file with 6 additions and 0 deletions.
diff --git a/molpipeline/estimators/chemprop/component_wrapper.py b/molpipeline/estimators/chemprop/component_wrapper.py
@@ -158,6 +158,12 @@ def __init__(
             Activation function.
         criterion : LossFunction or None, optional (default=None)
             Loss function. None defaults to BCELoss.
+        task_weights : Tensor or None, optional (default=None)
+            Task weights.
+        threshold : float or None, optional (default=None)
+            Threshold for binary classification.
+        output_transform : UnscaleTransform or None, optional (default=None)
+            Transformations to apply to the output. None defaults to UnscaleTransform.
         """
         if criterion is None:
             task_weights = torch.ones(n_tasks) if task_weights is None else task_weights