Save a dataset in a determined number of shards
This is useful to distribute the shards to training nodes.

This can be implemented in `save_to_disk` and can also leverage multiprocessing to speed up the process