Bug fix in get_pred_dict

AdvancedPhotonSource · May 31, 2024 · ae64f07 · ae64f07
1 parent 50658f3
commit ae64f07
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/generic_trainer/trainer.py b/generic_trainer/trainer.py
@@ -592,8 +592,8 @@ def get_pred_dict(self, preds):
         :return: dict.
         """
         d = {}
-        for i, name in enumerate(self.configs.pred_names_and_types):
-            d[name] = preds[i][0]
+        for i, name_and_type in enumerate(self.configs.pred_names_and_types):
+            d[name_and_type[0]] = preds[i]
         return d
 
     def process_data_loader_yield_sample_first(self, data_and_labels, data_label_separation_index):