Fixing restart bug