Add TensorBatch type

Author: markus-kreft

ocf_datapipes/batch/__init__.py

@@ -1,5 +1,5 @@
 """Datapipes for batching together data"""
-from .batches import BatchKey, NumpyBatch, NWPBatchKey, NWPNumpyBatch, XarrayBatch
+from .batches import BatchKey, NumpyBatch, NWPBatchKey, NWPNumpyBatch, TensorBatch, XarrayBatch
 from .merge_numpy_examples_to_batch import (
     MergeNumpyBatchIterDataPipe as MergeNumpyBatch,
 )

ocf_datapipes/batch/batches.py

@@ -4,6 +4,7 @@
 from typing import Union
 
 import numpy as np
+import torch
 import xarray as xr
 
 
@@ -229,3 +230,5 @@ class NWPBatchKey(Enum):
 NumpyBatch = dict[BatchKey, Union[np.ndarray, dict[str, NWPNumpyBatch]]]
 
 XarrayBatch = dict[BatchKey, Union[xr.DataArray, xr.Dataset]]
+
+TensorBatch = dict[BatchKey, Union[torch.Tensor, dict[str, dict[NWPBatchKey, torch.Tensor]]]]