diff --git a/utils/ipc_data_generator.py b/utils/ipc_data_generator.py index a821c9f..c03a642 100644 --- a/utils/ipc_data_generator.py +++ b/utils/ipc_data_generator.py @@ -1,5 +1,6 @@ import multiprocessing as mp from multiprocessing import shared_memory +import signal from typing import Callable, Optional, Tuple import numpy as np @@ -85,16 +86,14 @@ class IPCBatchGenerator: self.current_cache = 1 self.process_id = 'prefetch' + signal.signal(signal.SIGINT, signal.SIG_IGN) while self.prefetch_stop.buf is not None and self.prefetch_stop.buf[0] == 0: - try: - self.current_cache = 1 - self.current_cache - self.global_step += 1 + self.current_cache = 1 - self.current_cache + self.global_step += 1 - self._next_batch() - self.prefetch_pipe_child.recv() - self.prefetch_pipe_child.send(self.current_cache) - except KeyboardInterrupt: - break + self._next_batch() + self.prefetch_pipe_child.recv() + self.prefetch_pipe_child.send(self.current_cache) def _next_batch(self): # Loading data