Improve readability

pytorch · Feb 15, 2023 · eee76b3 · eee76b3
1 parent 09ffa9e
commit eee76b3
Showing 1 changed file with 11 additions and 10 deletions.
diff --git a/torchdata/datapipes/iter/util/converter.py b/torchdata/datapipes/iter/util/converter.py
@@ -86,17 +86,18 @@ def __getitem__(self, index):
             if self._map is None:
                 self._map = {}
                 self._itr = iter(self.datapipe)
-                raise KeyError
-            return self._map[index]
+            else:
+                return self._map[index]
         except KeyError:
-            while not self._depleted:
-                try:
-                    key, value = self._load_next_item()
-                    if key == index:
-                        return value
-                except StopIteration:
-                    self._depleted = True
-            raise IndexError(f"Index {index} is invalid for IterToMapConverter.")
+            pass
+        while not self._depleted:
+            try:
+                key, value = self._load_next_item()
+                if key == index:
+                    return value
+            except StopIteration:
+                self._depleted = True
+        raise IndexError(f"Index {index} is invalid for IterToMapConverter.")
 
     def _load_next_item(self):
         elem = next(self._itr)