Adding "state_counter" logging to average_single_index_loss_with_state.
Adding "state_counter" logging to average_single_index_loss_with_state.

PiperOrigin-RevId: 411094869
Change-Id: Idcc347454c6188950e2ffaf88c5f77e81b48d91d

deepmind/enn,deepmind/enn