Internal

PiperOrigin-RevId: 699267409
google · Nov 22, 2024 · 0f03abc · 0f03abc
1 parent 8954853
commit 0f03abc
Showing 1 changed file with 1 addition and 2 deletions.
diff --git a/grain/_src/python/dataset/transformations/prefetch.py b/grain/_src/python/dataset/transformations/prefetch.py
@@ -157,8 +157,7 @@ def __next__(self) -> T:
       if return_element:
         with self._stats.record_self_time(offset_ns=timer.value()):
           return element
-    with self._stats.record_self_time(offset_ns=timer.value()):
-      raise StopIteration
+    raise StopIteration
 
   def get_state(self):
     return {"next_index": self._next_index}