diff --git a/aim/ext/transport/request_queue.py b/aim/ext/transport/request_queue.py index 1301d0a60..d1242a50e 100644 --- a/aim/ext/transport/request_queue.py +++ b/aim/ext/transport/request_queue.py @@ -63,7 +63,7 @@ def worker(self): def _try_exec_task(self, task_f, *args): # temporary workaround for M1 build - from websockets.exceptions import ConnectionClosedError + from websockets.exceptions import WebSocketException retry = 0 while retry < self.retry_count: @@ -79,7 +79,7 @@ def _try_exec_task(self, task_f, *args): try: task_f(*args) return True - except ConnectionClosedError as e: + except (WebSocketException, TimeoutError) as e: self._needs_reconnect = True retry += 1