Example to save and reload a model

examples/liramlvq_tecator.py

@@ -37,3 +37,12 @@ if __name__ == "__main__":
 
     # Training loop
     trainer.fit(model, train_loader)
+
+    # Save the model
+    torch.save(model, "liramlvq_tecator.pt")
+
+    # Load a saved model
+    saved_model = torch.load("liramlvq_tecator.pt")
+
+    # Display the Lambda matrix
+    saved_model.show_lambda()

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):