hotfix: save model on epoch end
Former-commit-id: 8cd802c53a27fff273746dd9c77305fd22f250ce [formerly 3172f1b49c82ed8d3405488983156ae3f62e0598] Former-commit-id: cad3afc409deb9ffc905d1927383799850241e68
This commit is contained in:
parent
4ce22005cf
commit
382b3ff91a
|
@ -70,6 +70,10 @@ class ArtifactLog(Callback):
|
||||||
def on_fit_start(self, trainer, pl_module):
|
def on_fit_start(self, trainer, pl_module):
|
||||||
self.best = 1
|
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):
|
def on_validation_epoch_start(self, trainer, pl_module):
|
||||||
self.dices = []
|
self.dices = []
|
||||||
|
|
||||||
|
|
Loading…
Reference in a new issue