code fix after upgrade