def METHOD_NAME(batch):
    """Helper function to pad nested data within each batch elements."""
    padded_dict_batch = {}
    if isinstance(batch, dict):
        for key, value in batch.items():
            padded_dict_batch[key] = METHOD_NAME(value)
        return padded_dict_batch
    rank = len(batch.shape)
    assert rank > 0
    missing_count = self.padded_batch_size - self.get_real_batch_size(
        batch
    )
    padding = backend.stack(
        [[0, missing_count]] + [[0, 0]] * (rank - 1)
    )
    return tf.pad(batch, padding, "constant")