Minor modifications to fix display after game over