import tensorflow as tf

    def update_model(self, *grads):
        params = [self.w0, self.b0, self.w1, self.b1]
        grads = [tf.cast(grad, tf.float32) for grad in grads]
        with tf.name_scope('update'):
            update_op = tf.group(*[
                param.assign(param - grad * self.LEARNING_RATE)
