Fix key

openclimatefix · Nov 27, 2023 · c246be0 · c246be0
1 parent c7565a0
commit c246be0
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/pvnet/data/utils.py b/pvnet/data/utils.py
@@ -25,7 +25,7 @@ def batch_to_tensor(batch):
 
 def split_batches(batch):
     """Splits a single batch of data."""
-    n_samples = batch[BatchKey.sensor].shape[0]
+    n_samples = batch[BatchKey.gsp].shape[0]
     keys = list(batch.keys())
     examples = [{} for _ in range(n_samples)]
     for i in range(n_samples):