store to pytorch dataset feature

[bkmi]

It would be great to also create a "to pytorch dataset" feature where a zarr store would produce a dataset which does not access the zarr store. Rather it copies the data to a data structure of tensors.

The reason I want this is because swyft would be sped up significantly if we can use data already on the gpu, for example.