Small fix to handle stop state