Display plot during the training loop

app.py

@@ -59,6 +59,7 @@ fig, ax = plt.subplots(dpi=120)
 ax.set_xlim((0,1))
 ax.set_ylim((-5,26))
 ax.scatter(X[:,1], y, c='#e76254' ,edgecolors='firebrick')
+st.pyplot(fig)
 
 if cost_function == 'RMSE-Loss':
 
@@ -94,8 +95,10 @@ if cost_function == 'RMSE-Loss':
 
              # add a line to the plot
              plt.plot(X[jnp.dot(X, w).argsort(), 1], jnp.dot(X, w).sort(), 'c--')
+             st.pyplot(fig)
      # Plot the final line
      plt.plot(X[jnp.dot(X, w).argsort(), 1], jnp.dot(X, w).sort(), 'k-', label='Final line')
+     st.pyplot(fig)
      status_text.text('Done!')