Minor upgrade on code