small fix of restart bug

small fix of restart bug