def METHOD_NAME(self, dataset: Dataset) -> Dataset:
    if self.align_data:
        funcs = {"target": self._align_data_entry}
    if self.num_test_dates is None:
        grouped_dataset = self._prepare_train_data(dataset, funcs)
    else:
        grouped_dataset = self._prepare_test_data(dataset)
    return grouped_dataset