Alexandru Gherghescu commited on
Commit
3824b4f
1 Parent(s): 04fbb43

Save pre training losses at the end

Browse files
Files changed (1) hide show
  1. pre_training.py +3 -0
pre_training.py CHANGED
@@ -1,3 +1,4 @@
 
1
  from torch.optim import Adam
2
  from transformers import (
3
  AutoTokenizer,
@@ -75,4 +76,6 @@ print("Starting training...")
75
 
76
  trainer.train()
77
 
 
 
78
  trainer.save_model('trained')
 
1
+ import torch
2
  from torch.optim import Adam
3
  from transformers import (
4
  AutoTokenizer,
 
76
 
77
  trainer.train()
78
 
79
+ torch.save(trainer.state.log_history, 'trainer_history.pt')
80
+
81
  trainer.save_model('trained')