The dataset slicing should be implemented with some sort of DatasetView object which does not copy the data, but can provide all of the same utility functions and parameters. Maybe check and see how pytorch does this sort of thing and copy their approach?