Slow compared to simple loader

I'm trying a simple dataloader on MNIST and am finding it to be much slower than a simple dataloader I wrote to approximate the functionality (just permuting a range of indices + batching). Not sure if this is an issue or just related to my problem setup. Thanks in advance!

```
ds = tfds.data_source("mnist", split="train")
dl_grain = pygrain.DataLoader(
data_source=ds,
sampler=pygrain.IndexSampler(
num_records=len(ds),
shard_options=pygrain.NoSharding(),
shuffle=True,
num_epochs=1,
seed=123,
),
operations=[pygrain.Batch(batch_size=32)],
worker_count=0,
)
```