TypeError when applying map after set_format(type='torch')

[lhoestq]

Hi ! I think this is an issue in datasets, let me open a PR to fix this.

In the meantime you need to make sure every tensor in the returned batch dict should be lists or numpy arrays