Example to save and reload a model

2021-05-10 14:30:02 +02:00
parent ed03ab168e
commit 49100f43f5
2 changed files with 12 additions and 0 deletions
--- a/prototorch/models/vis.py
+++ b/prototorch/models/vis.py
@@ -356,6 +356,9 @@ class Vis2DAbstract(pl.Callback):
        else:
            plt.show(block=True)

+    def on_train_end(self, trainer, pl_module):
+        plt.show()
+

 class VisGLVQ2D(Vis2DAbstract):
    def on_epoch_end(self, trainer, pl_module):