[tfds] Fix tfds.as_numpy which now returns a reusable iterable.
[tfds] Fix tfds.as_numpy which now returns a reusable iterable.

`next(ds)` -> `next(iter(ds))`

PiperOrigin-RevId: 324977242
Change-Id: I8294fd695301321166edeaf949e5aeea0c97e640

deepmind/dm-haiku,deepmind/dm-haiku