min

run_mlm_flax_stream.py

@@ -564,7 +564,7 @@ if __name__ == "__main__":
     train_metrics = []
     eval_metrics = []
     
-    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))
+    training_iter = iter(torch.utils.data.DataLoader(tokenized_datasets.with_format("torch"), batch_size=1, shuffle=False, num_workers=min(33,dataset.n_shards), collate_fn=lambda x: x))
     
     max_seq_length = min(data_args.max_seq_length, tokenizer.model_max_length)
     eval_samples = advance_iter_and_group_samples(training_iter, data_args.num_eval_samples, max_seq_length)