change number of shards

run_mlm_flax_stream.py

@@ -581,7 +581,7 @@ if __name__ == "__main__":
             shuffle_seed += 1
             tokenized_datasets.set_epoch(shuffle_seed)
             
-            training_iter = iter(torch.utils.data.DataLoader(tokenized_datasets.with_format("torch"), batch_size=1, shuffle=False, num_workers=dataset.n_shards, collate_fn=lambda x: x))
+            training_iter = iter(torch.utils.data.DataLoader(tokenized_datasets.with_format("torch"), batch_size=1, shuffle=False, num_workers=max(33,dataset.n_shards), collate_fn=lambda x: x))
             
             eval_samples = advance_iter_and_group_samples(training_iter, data_args.num_eval_samples, max_seq_length)
             samples = advance_iter_and_group_samples(training_iter, train_batch_size, max_seq_length)