Run_qa.py for Tensorflow does not save the model in output directory

I am very new to the topic so sorry if the question is dumb.

I am trying to train a Q/A model with example script for Question Answering for Tensorflow.
I use the following command to run the script:

python run_qa.py \
--model_name_or_path distilbert-base-cased \
--output_dir output \
--dataset_name squad \
--max_train_samples 40 \
--max_eval_samples 20 \
--do_train \
--do_eval \

And seems that everything goes fine and without errors, but after the script finishes the training I cannot find the trained model in output directory, what I see there are these two file eval_nbest_predictions.json and eval_predictions.json .
So what am I doing wrong? How can I make the script save the trained model in the output directory?

1 Like

@Andranik I am facing the same outcome. Were you able to identify the root cause?

Hello there :wave:

I realized what was wrong, and you can open a PR to fix it as well.
The saving is done under this callback. The callback should’ve been instantiated and called inside model.fit() but it doesn’t. I’ll open a PR to fix that (but you can also use this quick fix).