add support for single class classification

ChEB-AI · Dec 14, 2023 · e90d66b · e90d66b
1 parent 462379b
commit e90d66b
Showing 1 changed file with 4 additions and 3 deletions.
diff --git a/chebai/result/classification.py b/chebai/result/classification.py
@@ -85,10 +85,11 @@ def evaluate_model(
                 save_ind += 1
                 n_saved = 0
 
-    test_preds = torch.cat(preds_list)
-    test_labels = torch.cat(labels_list)
+    if buffer_dir is None:
+        test_preds = torch.cat(preds_list)
+        test_labels = torch.cat(labels_list)
 
-    return test_preds, test_labels
+        return test_preds, test_labels
 
 
 def load_results_from_buffer(buffer_dir, device):