fix : gpu loss plotting

2023-11-29 16:31:28 +01:00
parent 42cf7fcaf1
commit 4639374961
1 changed files with 1 additions and 1 deletions
--- a/pina/plotter.py
+++ b/pina/plotter.py
@@ -256,7 +256,7 @@ class Plotter:
                )
            loss = trainer_metrics[metric]
            epochs = range(len(loss))
-            plt.plot(epochs, loss, **kwargs)
+            plt.plot(epochs, loss.cpu(), **kwargs)

        # plotting
        plt.xlabel('epoch')