fix that is valid if we assume that all datasets to merge have the same amount of tasks

        X, y, w, ids = (dataset.X, dataset.y, dataset.w, dataset.ids)

        yield (X, y, w, ids)

    return DiskDataset.create_dataset(generator(), data_dir=merge_dir)

    return DiskDataset.create_dataset(generator(), data_dir=merge_dir, tasks=next(iter(datasets)).tasks)

  def subset(self, shard_nums, subset_dir=None):

    """Creates a subset of the original dataset on disk."""