small code tweek