Use itertools instead of custom count

Itertools.count() does the same thing as
what the previous code does, so just use
that instead of custom code. This helps
reduce the amount of custom code.

Change-Id: I549889e9e1d1f607dfdbb467feffdd66dc825999