FIX: fix the crash when dataset object returns lists instead of tuples

AdvancedPhotonSource · Dec 5, 2024 · d0fdaeb · d0fdaeb
1 parent 76a2940
commit d0fdaeb
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/generic_trainer/trainer.py b/generic_trainer/trainer.py
@@ -646,7 +646,7 @@ def process_data_loader_yield(self, data_and_labels: Any, data_label_separation_
                                             to be data.
         :returns tuple[torch.tensor], tuple[torch.tensor]. 2 tuples for data and labels.
         """
-        if isinstance(data_and_labels[0], tuple):
+        if isinstance(data_and_labels[0], (tuple, list)):
             # In this case, data_and_labels is in sample-then-item order.
             data_list, label_list = self.process_data_loader_yield_sample_first(data_and_labels,
                                                                                 data_label_separation_index)