change number of shards

events.out.tfevents.1672853660.t1v-n-29919176-w-3.257607.0.v2

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:7e5974ad314fef8b8325e605291a232c7513c39cb33095f5bcd7fd49c8867935
-size 2577219
+oid sha256:df7888d10faf8f207ed12292b63244b752f65670bf6880c4d29a0717ede95467
+size 2651789

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)