diff --git a/mindspore/dataset/engine/iterators.py b/mindspore/dataset/engine/iterators.py index 343fe8dd07..c71f3c3b4e 100644 --- a/mindspore/dataset/engine/iterators.py +++ b/mindspore/dataset/engine/iterators.py @@ -15,6 +15,8 @@ """Built-in iterators. """ from abc import abstractmethod +import os +import signal import weakref import numpy as np @@ -160,7 +162,16 @@ class DictIterator(Iterator): Returns: Dict, the next record in the dataset. """ - return {k: self._transform_tensor(t) for k, t in self._iterator.GetNextAsMap().items()} + try: + return {k: self._transform_tensor(t) for k, t in self._iterator.GetNextAsMap().items()} + except RuntimeError as err: + ## maybe "Out of memory" / "MemoryError" error + logger.error("Got runtime err: {}.".format(err)) + err_info = str(err) + if err_info.find("Out of memory") >= 0 or err_info.find("MemoryError") >= 0: + logger.error("Memory error occurred, process will exit.") + os.kill(os.getpid(), signal.SIGKILL) + raise err class TupleIterator(Iterator):