lint: docstrings and tests

basf · Sep 3, 2024 · fdb1d31 · fdb1d31
1 parent 4a117d5
commit fdb1d31
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 4 deletions.
diff --git a/molpipeline/estimators/chemprop/component_wrapper.py b/molpipeline/estimators/chemprop/component_wrapper.py
@@ -189,6 +189,8 @@ def __init__(
             Threshold for binary classification.
         output_transform : UnscaleTransform or None, optional (default=None)
             Transformations to apply to the output. None defaults to UnscaleTransform.
+        kwargs : Any
+            Additional keyword arguments.
         """
         if task_weights is None:
             task_weights = torch.ones(n_tasks)

diff --git a/test_extras/test_chemprop/test_models.py b/test_extras/test_chemprop/test_models.py
@@ -13,7 +13,6 @@
     BondMessagePassing,
     MeanAggregation,
     MulticlassClassificationFFN,
-    BinaryClassificationFFN,
     RegressionFFN,
     SumAggregation,
 )
@@ -251,12 +250,12 @@ def test_error_for_multiclass_predictor(self) -> None:
         with self.assertRaises(ValueError):
             predictor = MulticlassClassificationFFN(n_classes=2)
             model = MPNN(message_passing=bond_encoder, agg=agg, predictor=predictor)
-            chemprop_model = ChempropMulticlassClassifier(n_classes=2, model=model)
+            ChempropMulticlassClassifier(n_classes=2, model=model)
         with self.assertRaises(ValueError):
             predictor = MulticlassClassificationFFN(n_classes=3)
             model = MPNN(message_passing=bond_encoder, agg=agg, predictor=predictor)
-            chemprop_model = ChempropMulticlassClassifier(n_classes=4, model=model)
+            ChempropMulticlassClassifier(n_classes=4, model=model)
         with self.assertRaises(AttributeError):
             predictor = RegressionFFN()
             model = MPNN(message_passing=bond_encoder, agg=agg, predictor=predictor)
-            chemprop_model = ChempropMulticlassClassifier(n_classes=4, model=model)
+            ChempropMulticlassClassifier(n_classes=4, model=model)