Fixrefactor

fix some bug introduced by the refactor