ENSEEIHT/projet-long
1
0
Fork 0
Code Issues Pull requests Packages Projects Releases Wiki Activity

hotfix: save model on epoch end

Browse source 
Former-commit-id: 8cd802c53a27fff273746dd9c77305fd22f250ce [formerly 3172f1b49c82ed8d3405488983156ae3f62e0598]
Former-commit-id: cad3afc409deb9ffc905d1927383799850241e68
This commit is contained in:
Laurent Fainsin 2022-08-26 11:10:07 +02:00
parent 4ce22005cf
commit 382b3ff91a
1 changed files with 4 additions and 0 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

4
src/utils/callback.py
View file

@ -70,6 +70,10 @@ class ArtifactLog(Callback):
def on_fit_start(self, trainer, pl_module):
self.best = 1
def on_train_epoch_end(self, trainer, pl_module):
# create checkpoint
torch.save(pl_module.state_dict(), "checkpoints/model.pth")
def on_validation_epoch_start(self, trainer, pl_module):
self.dices = []