fix : gpu loss plotting

mathLab · Nov 29, 2023 · 4639374 · 4639374
1 parent 42cf7fc
commit 4639374
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/pina/plotter.py b/pina/plotter.py
@@ -256,7 +256,7 @@ def plot_loss(self,
                 )
             loss = trainer_metrics[metric]
             epochs = range(len(loss))
-            plt.plot(epochs, loss, **kwargs)
+            plt.plot(epochs, loss.cpu(), **kwargs)
 
         # plotting
         plt.xlabel('epoch')