From 382b3ff91a6cb4c9508f537652f1fc2dcfebfc77 Mon Sep 17 00:00:00 2001
From: Laurent Fainsin <laurentfainsin@protonmail.com>
Date: Fri, 26 Aug 2022 11:10:07 +0200
Subject: [PATCH] hotfix: save model on epoch end

Former-commit-id: 8cd802c53a27fff273746dd9c77305fd22f250ce [formerly 3172f1b49c82ed8d3405488983156ae3f62e0598]
Former-commit-id: cad3afc409deb9ffc905d1927383799850241e68
---
 src/utils/callback.py | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/src/utils/callback.py b/src/utils/callback.py
index 15cd5e5..7adf0e0 100644
--- a/src/utils/callback.py
+++ b/src/utils/callback.py
@@ -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 = []