From e90d66b501fd3e9eb988edd3e8b5b1868bca5d4a Mon Sep 17 00:00:00 2001 From: sfluegel Date: Thu, 14 Dec 2023 15:39:24 +0100 Subject: [PATCH] add support for single class classification --- chebai/result/classification.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/chebai/result/classification.py b/chebai/result/classification.py index f1191e16..05ada33b 100644 --- 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):