Fixed reshape error for invalid smiles (#229)

* fixed reshape error for invalid smiles * fixed parallel issue with ind_preds variable
chemprop · Jan 21, 2022 · 1e264cc · 1e264cc
1 parent cf420db
commit 1e264cc
Showing 1 changed file with 4 additions and 3 deletions.
diff --git a/chemprop/train/make_predictions.py b/chemprop/train/make_predictions.py
@@ -206,9 +206,10 @@ def predict_and_save(args: PredictArgs, train_args: TrainArgs, test_data: Molecu
 
         # Reshape multiclass to merge task and class dimension, with updated num_tasks
         if args.dataset_type == 'multiclass':
-            preds = preds.reshape((num_tasks))
-            if args.ensemble_variance or args. individual_ensemble_predictions:
-                ind_preds = ind_preds.reshape((num_tasks, len(args.checkpoint_paths)))
+            if isinstance(preds, np.ndarray) and preds.ndim > 1:
+                preds = preds.reshape((num_tasks))
+                if args.ensemble_variance or args. individual_ensemble_predictions:
+                    ind_preds = ind_preds.reshape((num_tasks, len(args.checkpoint_paths)))
 
         # If extra columns have been dropped, add back in SMILES columns
         if args.drop_extra_columns: