Fix key

openclimatefix · Feb 6, 2024 · 30908f1 · 30908f1
1 parent 809ebbf
commit 30908f1
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/pvnet/models/utils.py b/pvnet/models/utils.py
@@ -119,7 +119,7 @@ def flush(self) -> dict[BatchKey, list[torch.Tensor]]:
         batch = {}
         for k, v in self._batches.items():
             print(k)
-            if k == f"{self.key_to_keep}_t0_idx":
+            if k == BatchKey[f"{self.key_to_keep}_t0_idx"]:
                 batch[k] = v[0]
             else:
                 batch[k] = torch.cat(v, dim=0)