Fix some issue caused by refactoring.