Fixed code after merge