Add the schema to the output of the .repartition() method

[sararb]

• For multi-gpu training support in transformers4rec, we are using .repartition() to split the dataset across the number of available GPUs (i.e. global_size) equally. (in this line).
• The multi-gpu script is broken because the repartition method returns a new Dataset object without copying the original dataset schema (where we set properties such as is_ragged=False and value_count to ensure the dataloader returns dense tensors instead of a tuple representation).
• This PR proposes a quick-fix where the original schema is passed to the new Dataset object returned by dataset.repartition()

[github-actions]

Documentation preview

https://nvidia-merlin.github.io/core/review/pr-192