Fix "schedule while atomic" error