Update code.

Signed-off-by: Geidivan <geidivan@gmail.com>