Convert dataset to pytorch dataloader

Also, for whatever reason, it’s not caching the .map() transform, even though the hash for the lambda is the same on every session. What’s going on there?