Use isLeafTask() to check whether to save state.

We meant to save states for all leaf tasks, including one layer task,
which is excluded by the condition checked before.

Bug: 152318088
Test: Smoke tests.
Change-Id: Iadd8aa60ffe254398d242a12dae39148314891c8