From c246be0cc49b80f06b542fe32049142fb8c0d0cb Mon Sep 17 00:00:00 2001
From: Jacob Bieker <jacob@bieker.tech>
Date: Mon, 27 Nov 2023 12:01:32 +0000
Subject: [PATCH] Fix key

---
 pvnet/data/utils.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/pvnet/data/utils.py b/pvnet/data/utils.py
index cd77c887..c85ce40c 100644
--- 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):