Optimize the code


Signed-off-by: jssoscar