fix the issues due to refactoring