import numpy as np
from autograd import grad, elementwise_grad

N = 5
x = np.linspace(0, 1, N)

n_neurons = 5

w_h = np.random.randn(n_neurons, N)
w_o = np.random.randn(n_neurons, 1)
w = [w_h, w_o]

def sigmoid(z):
    return 1/(1+np.exp(-z))

def NN(x, w):
    a = np.dot(w_h, x)
    a1 = sigmoid(a)
    b = np.dot(a1, w_o)
    b1 = sigmoid(b)
    return b1 

def gt(x, w):
    g0 = 1
    return g0+x*NN(x, w)

def g(x, gt):
    return x**3 + 2.*x + x**2 * ((1. + 3.*x**2) / (1. + x + x**3))-gt(x, w)*(x + (1. + 3.*x**2) / (1. + x + x**3))

def d_segmoid(z):
    return sigmoid(z)*(1-sigmoid(z))

def d_NN(x, w):
    k = 1
    a11 = np.dot(w_o.T, w_h.T**k)
    return np.dot(a11, d_segmoid(x))


def loss_function(x, w):
    
    err = (d_NN(x, w)-g(x, gt))**2
    cost_sum = np.sum(err)
    
    return cost_sum/ np.size(err)

def optimalw(x, lmb):
    lmb = 0.001
    for i in range(1000):
        loss_function_grad = grad(loss_function,0)
    
        w[0] = w[0]-lmb*loss_function_grad[0]
        w[1] = w[1]-lmb*loss_function_grad[1]
    
    return w

Error:
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
<ipython-input-26-9381da78f871> in <module>
      8 
      9     return w
---> 10 optimalw(x, lmb)

<ipython-input-26-9381da78f871> in optimalw(x, lmb)
      4         loss_function_grad = grad(loss_function,0)
      5 
----> 6         w[0] = w[0]-lmb*loss_function_grad[0]
      7         w[1] = w[1]-lmb*loss_function_grad[1]
      8 

TypeError: 'function' object is not subscriptable