Make with_format correctly indicate that a Dataset is compatible with PyTorch's Dataset class

[benglewis]

Feature request

Currently datasets does not correctly indicate to the Python type-checker (e.g. pyright / Pylance) that the output of with_format is compatible with PyTorch's Dataloader since it does not indicate that the HuggingFace Dataset is compatible with the PyTorch Dataset class. It would be great if we could get the typing to work nicely.

Motivation

To avoid casting types in our Python code.

Your contribution

I would be happy to contribute a PR if this is something that may be accepted and could work with the current approach.
This doesn't have to be for just PyTorch, but I imagine that the same thing would be useful for tensorflow and such, but we only have a need for PyTorch at this stage.