Skip to content
Snippets Groups Projects
Commit efdb2601 authored by friebolin's avatar friebolin
Browse files

Save models

parent 5a5875fc
No related branches found
No related tags found
No related merge requests found
......@@ -190,7 +190,7 @@ def train(model, name, imdb, seed,gradient_accumulation_steps,mixup, threshold,
#progress_bar.update(1)
#print("one epoch done")
torch.save(model.state_dict(), "bert_baseline.pt")
torch.save(model, "./saved_models/bert_baseline.pt")
#print(model_name)
evaluation_test = evaluation.evaluate_model(model, name, test_dataset, learning_rate, test_batch_size, imdb)
......
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment