python torch_dataset_loader.py --shards_index pretraining_data/shards/pretraining_dataset.index.json \
  --batch_size 32 --num_workers 4 --shuffle_buffer 10000 --repeat 1 --max_batches 5