fix param_count in build_model

Author: AnFreTh

mambular/models/sklearn_base_lss.py

@@ -30,6 +30,7 @@
     NormalDistribution,
     PoissonDistribution,
     StudentTDistribution,
+    Quantile,
 )
 from lightning.pytorch.callbacks import ModelSummary
 
@@ -210,11 +211,9 @@ def build_model(
             X, y, X_val, y_val, val_size=val_size, random_state=random_state
         )
 
-        num_classes = len(np.unique(y))
-
         self.task_model = TaskModel(
             model_class=self.base_model,
-            num_classes=num_classes,
+            num_classes=self.family.param_count,
             config=self.config,
             cat_feature_info=self.data_module.cat_feature_info,
             num_feature_info=self.data_module.num_feature_info,
@@ -347,6 +346,7 @@ def fit(
             "negativebinom": NegativeBinomialDistribution,
             "inversegamma": InverseGammaDistribution,
             "categorical": CategoricalDistribution,
+            "quantile": Quantile,
         }
 
         if distributional_kwargs is None: