pass correct tasks

basf · Sep 3, 2024 · a064100 · a064100
1 parent c5810fc
commit a064100
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/molpipeline/estimators/chemprop/component_wrapper.py b/molpipeline/estimators/chemprop/component_wrapper.py
@@ -340,7 +340,7 @@ def __init__(
         output_transform: UnscaleTransform | None = None,
     ):
         super().__init__(
-            n_tasks * n_classes,
+            n_tasks,
             input_dim,
             hidden_dim,
             n_layers,