Fix issues after upgrade