Fix code

app.py

@@ -49,18 +49,22 @@ st.markdown('Using gradient descent we find the minima of the loss adjusting the
 
 st.latex(r'''\bf{w}\leftarrow \bf{w}-\eta \frac{\partial\ell(\bf{X},\bf{y}, \bf{w})}{\partial \bf{w}}''')
 
+st.markdown('The training loop:')
+
+code = '''NUM_ITER = 1000
+# initialize parameters
+w = np.array([3., -2., -8.])
+for i in range(NUM_ITER):
+     # update parameters
+     w -= learning_rte * grad_loss(w)'''
+     
+st.code(code, language='python')
 
 # Fitting by the respective cost_function
 w = jnp.array(np.random.random(3))
 learning_rate = 0.05
 NUM_ITER = 1000
 
-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':
 
      def loss(w):
@@ -92,12 +96,12 @@ if cost_function == 'RMSE-Loss':
              # report the loss at the current epoch
              status_text.text(
                  'Trained loss at epoch %s is %s' % (i, loss(w)))
-
-             # 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')
+     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')
+     ax.plot(X[jnp.dot(X, w).argsort(), 1], jnp.dot(X, w).sort(), 'k-', label='Final line')
      st.pyplot(fig)
      status_text.text('Done!')
 
@@ -110,14 +114,3 @@ else:
       (y^{(i)}-\hat{y}^{(i)})^2 & \text{for }\quad |y^{(i)}-\hat{y}^{(i)}|\leq \delta \\
       2\delta|y^{(i)}-\hat{y}^{(i)}| - \delta^2 & \text{otherwise}
 \end{cases}''')
-
-st.markdown('The training loop:')
-
-code = '''NUM_ITER = 1000
-# initialize parameters
-w = np.array([3., -2., -8.])
-for i in range(NUM_ITER):
-     # update parameters
-     w -= learning_rte * grad_loss(w)'''
-     
-st.code(code, language='python')