diff --git a/test/test_lm_finetuning.py b/test/test_lm_finetuning.py index 6052c7cd6..b0d3684c9 100644 --- a/test/test_lm_finetuning.py +++ b/test/test_lm_finetuning.py @@ -36,7 +36,7 @@ def test_lm_finetuning(caplog): max_seq_len=12, next_sent_pred=True ) - data_silo = DataSilo(processor=processor, batch_size=batch_size) + data_silo = DataSilo(processor=processor, batch_size=batch_size, max_processes=1) language_model = LanguageModel.load(lang_model) lm_prediction_head = BertLMHead.load(lang_model) @@ -113,7 +113,7 @@ def test_lm_finetuning_no_next_sentence(caplog): max_seq_len=12, next_sent_pred=False ) - data_silo = DataSilo(processor=processor, batch_size=batch_size) + data_silo = DataSilo(processor=processor, batch_size=batch_size, max_processes=1) language_model = LanguageModel.load(lang_model) lm_prediction_head = BertLMHead.load(lang_model) @@ -190,7 +190,7 @@ def test_lm_finetuning_custom_vocab(caplog): max_seq_len=12, next_sent_pred=True ) - data_silo = DataSilo(processor=processor, batch_size=batch_size) + data_silo = DataSilo(processor=processor, batch_size=batch_size, max_processes=1) language_model = LanguageModel.load(lang_model, n_added_tokens=len(tokenizer.added_tokens_decoder)) lm_prediction_head = BertLMHead.load(lang_model, n_added_tokens=len(tokenizer.added_tokens_decoder))