Return on the end of the prefetch_to_device generator.
Return on the end of the prefetch_to_device generator.

Following pep-0479, modern python generators signal the end by returning rather than raising StopIteration.
(this issue was originally noted by Marvin Ritter.)
ref: https://www.python.org/dev/peps/pep-0479/
google/flax,google/flax