fixing keys in dataloader _split_condition

Author: dario-coscia

pina/data/data_module.py

@@ -427,11 +427,21 @@ def _split_condition(single_condition_dict, splits_dict):
         offset = 0
 
         for stage, stage_len in splits_dict.items():
+            # filter keys to return dict, only batch items are listed here
+            # Equations are NEVER dataloaded, None variables are not dataloaded
+            filtered_keys = []
+            for k in single_condition_dict.keys():
+                if k == "equation" or (
+                    k == "conditional_variables"
+                    and single_condition_dict[k] is None
+                ):
+                    continue
+                else:
+                    filtered_keys.append(k)
+            # return the correct splitting
             to_return_dict[stage] = {
-                k: v[offset : offset + stage_len]
-                for k, v in single_condition_dict.items()
-                if k != "equation"
-                # Equations are NEVER dataloaded
+                k: single_condition_dict[k][offset : offset + stage_len]
+                for k in filtered_keys
             }
             if offset + stage_len >= len_condition:
                 offset = len_condition - 1
@@ -456,6 +466,8 @@ def _apply_shuffle(condition_dict, len_data):
             for k, v in condition_dict.items():
                 if k == "equation":
                     continue
+                if k == "conditional_variables" and condition_dict[k] is None:
+                    continue
                 if isinstance(v, list):
                     condition_dict[k] = [v[i] for i in idx]
                 elif isinstance(v, LabelTensor):