Code - fix bug