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

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

PiperOrigin-RevId: 324977229
Change-Id: I17506d76a6b99eb8b64f43a88e94d57f12e23598

examples/mnist.py

@@ -51,7 +51,7 @@ def load_dataset(
   if is_training:
     ds = ds.shuffle(10 * batch_size, seed=0)
   ds = ds.batch(batch_size)
-  return tfds.as_numpy(ds)
+  return iter(tfds.as_numpy(ds))
 
 
 def main(_):