[FIX] for predictions with dataset.n_tasks = 0

Author: jcapels

setup.cfg

@@ -1,6 +1,6 @@
 [metadata]
 name = deepmol
-version = 1.1.16
+version = 1.1.17
 description = DeepMol: a python-based machine and deep learning framework for drug discovery
 keywords = machine-learning, deep-learning, cheminformatics, drug-discovery
 author = DeepMol Team

src/deepmol/__init__.py

@@ -1,2 +1,2 @@
 
-__version__ = '1.1.16'
+__version__ = '1.1.17'

src/deepmol/models/keras_models.py

@@ -214,11 +214,11 @@ def predict_proba(self, dataset: Dataset, return_invalid: bool = False) -> np.nd
                 self.logger.info(str(self.model))
                 self.logger.info(str(type(self.model)))
                 predictions = self.model.predict(dataset.X.astype('float32'))
-
+        
         predictions = np.array(predictions)
         if predictions.shape != (len(dataset.mols), dataset.n_tasks):
             predictions = normalize_labels_shape(predictions, dataset.n_tasks)
-
+        
         if len(predictions.shape) > 1:
             if predictions.shape[1] == len(dataset.mols) and predictions.shape[0] == dataset.n_tasks:
                 predictions = predictions.T

src/deepmol/utils/utils.py

@@ -167,8 +167,10 @@ def normalize_labels_shape(y_pred: Union[List, np.ndarray], n_tasks: int) -> np.
     """
     if not isinstance(y_pred, np.ndarray):
         y_pred = np.array(y_pred)
-
-    if n_tasks == 1:
+    
+    if n_tasks == 0 and len(y_pred.shape) < 2:
+        labels = _normalize_singletask_labels_shape(y_pred)
+    elif n_tasks == 1:
         labels = _normalize_singletask_labels_shape(y_pred)
     else:
         if len(y_pred.shape) == 3: