added confusion plots in plot_data

ThomasHelfer · Jul 1, 2024 · 022ab76 · 022ab76
1 parent 3f3b03a
commit 022ab76
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/evaluate_models.py b/evaluate_models.py
@@ -177,15 +177,15 @@
     model = model.to(device)
     model.eval()
 
-    y_true, y_true_label, y_pred = process_data_loader(
+    y_true, y_true_label, y_pred, lc_data = process_data_loader(
         val_loader_no_aug,
         regression,
         classification,
         device,
         model,
         combinations=cfg_extra_args["combinations"],
     )
-    y_true_train, y_true_train_label, _ = process_data_loader(
+    y_true_train, y_true_train_label, _, _ = process_data_loader(
         train_loader_no_aug,
         regression,
         classification,