def build_header(dataset):
    """Build header dictionary of metadata for the tensors in the dataset. This will be used when loading
    the tfrecords file to reconstruct the original tensors from the raw data. Shape is stored as an array
    and dtype is stored as an enumerated value (defined by tensorflow)."""
    header = {}
    for key in dataset.element_spec.keys():
        header[key] = {
            "shape": list(dataset.element_spec[key].shape),
            "dtype": dataset.element_spec[key].dtype.as_datatype_enum,
        }

    return header