High test loss for nonlinear ODE case

[edoardo100]

I modified both main() and ode_system() as stated in the instructions, but I get a very high test loss. Do I have to modify other functions? I also increased the number of epochs to 10^5 as pointed in the supplementary information of the cited paper (a very interesting work I have to say).
Thank you in advance.

Additional information:
Modified lines (all in src/deeponet_pde.py):
line 104: uncommented after having commented line 102
line 231: be sure to have problem="ode"
line 260: increased epochs to 100000

Example of output:

Step      Train loss    Test loss     Test metric   
0         [3.68e-01]    [1.74e+25]    [3.55e-01]    
1000      [3.08e-03]    [1.74e+25]    [2.05e-03]    
2000      [8.55e-04]    [1.74e+25]    [8.00e-04]    
3000      [3.57e-04]    [1.74e+25]    [5.24e-04]    
4000      [2.06e-04]    [1.74e+25]    [3.58e-04]    
5000      [1.65e-04]    [1.74e+25]    [3.22e-04]    
6000      [1.43e-04]    [1.74e+25]    [2.83e-04]    
7000      [1.28e-04]    [1.74e+25]    [2.47e-04]    
8000      [1.01e-04]    [1.74e+25]    [2.26e-04]    
9000      [8.53e-05]    [1.74e+25]    [2.00e-04]    
10000     [7.95e-05]    [1.74e+25]    [1.87e-04]    
11000     [7.90e-05]    [1.74e+25]    [1.80e-04]    
12000     [9.09e-05]    [1.74e+25]    [1.88e-04]    
13000     [9.10e-05]    [1.74e+25]    [1.69e-04]    
14000     [5.61e-05]    [1.74e+25]    [1.45e-04]    
15000     [5.24e-05]    [1.74e+25]    [1.40e-04]    
16000     [5.50e-05]    [1.74e+25]    [1.44e-04]    
17000     [5.00e-05]    [1.74e+25]    [1.33e-04]    
18000     [5.21e-05]    [1.74e+25]    [1.39e-04]    
19000     [4.54e-05]    [1.74e+25]    [1.31e-04]    
20000     [1.13e-04]    [1.74e+25]    [1.78e-04]    
21000     [4.39e-05]    [1.74e+25]    [1.31e-04]    
22000     [5.68e-05]    [1.74e+25]    [1.45e-04]    
23000     [4.37e-05]    [1.74e+25]    [1.31e-04]    
24000     [4.18e-05]    [1.74e+25]    [1.28e-04]    
25000     [4.00e-05]    [1.74e+25]    [1.28e-04]    
26000     [6.98e-05]    [1.74e+25]    [1.55e-04]    
27000     [7.45e-05]    [1.74e+25]    [1.58e-04]    
28000     [4.10e-05]    [1.74e+25]    [1.22e-04]    
29000     [3.62e-05]    [1.74e+25]    [1.20e-04]    
30000     [5.16e-05]    [1.74e+25]    [1.36e-04]    
31000     [3.60e-05]    [1.74e+25]    [1.19e-04]    
32000     [3.82e-05]    [1.74e+25]    [1.20e-04]    
33000     [3.33e-05]    [1.74e+25]    [1.14e-04]    
34000     [3.36e-05]    [1.74e+25]    [1.14e-04]    
35000     [3.24e-05]    [1.74e+25]    [1.11e-04]    
36000     [3.15e-05]    [1.74e+25]    [1.12e-04]    
37000     [3.14e-05]    [1.74e+25]    [1.08e-04]    
38000     [3.78e-05]    [1.74e+25]    [1.18e-04]    
39000     [3.40e-05]    [1.74e+25]    [1.09e-04]    
40000     [4.34e-05]    [1.74e+25]    [1.14e-04]    
41000     [3.58e-05]    [1.74e+25]    [1.09e-04]    
42000     [2.89e-05]    [1.74e+25]    [1.04e-04]    
43000     [2.99e-05]    [1.74e+25]    [1.09e-04]    
44000     [4.13e-05]    [1.74e+25]    [1.10e-04]    
45000     [2.72e-05]    [1.74e+25]    [1.03e-04]    
46000     [2.71e-05]    [1.74e+25]    [1.03e-04]    
47000     [2.88e-05]    [1.74e+25]    [1.02e-04]    
48000     [5.22e-05]    [1.74e+25]    [1.33e-04]    
49000     [2.55e-05]    [1.74e+25]    [9.98e-05]    
50000     [2.92e-05]    [1.74e+25]    [1.01e-04]    
51000     [2.63e-05]    [1.74e+25]    [9.88e-05]    
52000     [2.39e-05]    [1.74e+25]    [9.90e-05]    
53000     [2.76e-05]    [1.74e+25]    [1.02e-04]    
54000     [3.15e-05]    [1.74e+25]    [9.99e-05]    
55000     [3.11e-05]    [1.74e+25]    [1.11e-04]    
56000     [2.47e-05]    [1.74e+25]    [9.72e-05]    
57000     [3.14e-05]    [1.74e+25]    [1.08e-04]    
58000     [3.34e-05]    [1.74e+25]    [1.01e-04]    
59000     [2.54e-05]    [1.74e+25]    [9.68e-05]    
60000     [3.11e-05]    [1.74e+25]    [9.88e-05]    
61000     [4.05e-05]    [1.74e+25]    [1.20e-04]    
62000     [4.36e-05]    [1.74e+25]    [1.22e-04]    
63000     [2.21e-05]    [1.74e+25]    [9.88e-05]    
64000     [2.31e-05]    [1.74e+25]    [9.72e-05]    
65000     [3.04e-05]    [1.74e+25]    [9.75e-05]    
66000     [2.57e-05]    [1.74e+25]    [9.81e-05]    
67000     [2.77e-05]    [1.74e+25]    [9.60e-05]    
68000     [2.22e-05]    [1.74e+25]    [9.49e-05]    
69000     [2.21e-05]    [1.74e+25]    [9.86e-05]    
70000     [2.57e-05]    [1.74e+25]    [9.89e-05]    
71000     [2.21e-05]    [1.74e+25]    [9.48e-05]    
72000     [2.11e-05]    [1.74e+25]    [9.38e-05]    
73000     [2.57e-05]    [1.74e+25]    [1.02e-04]    
74000     [3.00e-05]    [1.74e+25]    [1.09e-04]    
75000     [2.15e-05]    [1.74e+25]    [9.34e-05]    
76000     [2.06e-05]    [1.74e+25]    [9.62e-05]    
77000     [3.68e-05]    [1.74e+25]    [1.16e-04]    
78000     [1.95e-05]    [1.74e+25]    [9.37e-05]    
79000     [2.38e-05]    [1.74e+25]    [9.98e-05]    
80000     [6.12e-05]    [1.74e+25]    [1.42e-04]    
81000     [2.97e-05]    [1.74e+25]    [1.08e-04]    
82000     [2.12e-05]    [1.74e+25]    [9.13e-05]    
83000     [2.11e-05]    [1.74e+25]    [9.13e-05]    
84000     [1.92e-05]    [1.74e+25]    [9.08e-05]    
85000     [4.75e-05]    [1.74e+25]    [1.28e-04]    
86000     [6.35e-05]    [1.74e+25]    [1.44e-04]    
87000     [1.88e-05]    [1.74e+25]    [9.50e-05]    
88000     [1.85e-05]    [1.74e+25]    [9.09e-05]    
89000     [1.74e-05]    [1.74e+25]    [9.12e-05]    
90000     [1.75e-05]    [1.74e+25]    [9.28e-05]    
91000     [3.79e-05]    [1.74e+25]    [1.05e-04]    
92000     [1.98e-05]    [1.74e+25]    [9.06e-05]    
93000     [2.04e-05]    [1.74e+25]    [9.09e-05]    
94000     [2.14e-05]    [1.74e+25]    [9.92e-05]    
95000     [1.70e-05]    [1.74e+25]    [9.22e-05]    
96000     [3.40e-05]    [1.74e+25]    [9.47e-05]    
97000     [1.80e-05]    [1.74e+25]    [9.49e-05]    
98000     [1.65e-05]    [1.74e+25]    [9.20e-05]    
99000     [1.67e-05]    [1.74e+25]    [8.97e-05]    
100000    [1.65e-05]    [1.74e+25]    [9.20e-05]    

Best model at step 98000:
  train loss: 1.65e-05
  test loss: 1.74e+25
  test metric: [9.20e-05]

[lululxvi]

See FAQ "Q: I failed to train the network or get the right solution, e.g., large training loss, unbalanced losses."