Modify plot_loss_history to handle irregular array sizes (#1571)

lululxvi · Nov 25, 2023 · c55e5d3 · c55e5d3
1 parent 26860b1
commit c55e5d3
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/deepxde/utils/external.py b/deepxde/utils/external.py
@@ -192,8 +192,10 @@ def plot_loss_history(loss_history, fname=None):
         fname (string): If `fname` is a string (e.g., 'loss_history.png'), then save the
             figure to the file of the file name `fname`.
     """
-    loss_train = np.sum(loss_history.loss_train, axis=1)
-    loss_test = np.sum(loss_history.loss_test, axis=1)
+    # np.sum(loss_history.loss_train, axis=1) is error-prone for arrays of varying lengths.
+    # Handle irregular array sizes.
+    loss_train = np.array([np.sum(loss) for loss in loss_history.loss_train])
+    loss_test = np.array([np.sum(loss) for loss in loss_history.loss_test])
 
     plt.figure()
     plt.semilogy(loss_history.steps, loss_train, label="Train loss")