import tensorflow as tf

    def net_U0(self, x):
        lambda_1 = self.lambda_1
        lambda_2 = tf.exp(self.lambda_2)
        U = self.neural_net(x, self.weights, self.biases)        
        U_x = self.fwd_gradients_0(U, x)
        U_xx = self.fwd_gradients_0(U_x, x)
        F = -lambda_1*U*U_x + lambda_2*U_xx
        U0 = U - self.dt*tf.matmul(F, self.IRK_alpha.T)
        return U0
    
    def net_U1(self, x):
        lambda_1 = self.lambda_1
        lambda_2 = tf.exp(self.lambda_2)
        U = self.neural_net(x, self.weights, self.biases)        
        U_x = self.fwd_gradients_1(U, x)
        U_xx = self.fwd_gradients_1(U_x, x)
        F = -lambda_1*U*U_x + lambda_2*U_xx
        U1 = U + self.dt*tf.matmul(F, (self.IRK_beta - self.IRK_alpha).T)
        return U1

    def callback(self, loss):
        print('Loss:', loss)
    
    def train(self, nIter):
