diff --git a/CHANGELOG.md b/CHANGELOG.md
index 1619be46..6069a6b9 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -12,6 +12,12 @@ Changes are grouped as follows
 - `Fixed` for any bug fixes.
 - `Security` in case of vulnerabilities.
 
+## [6.0.2]
+
+### Fixed
+
+ * Improved the state store retry behavior to handle connection errors
+
 ## [6.0.1]
 
 ### Fixed
diff --git a/cognite/extractorutils/__init__.py b/cognite/extractorutils/__init__.py
index f4092f1c..45ef551c 100644
--- a/cognite/extractorutils/__init__.py
+++ b/cognite/extractorutils/__init__.py
@@ -16,5 +16,5 @@
 Cognite extractor utils is a Python package that simplifies the development of new extractors.
 """
 
-__version__ = "6.0.1"
+__version__ = "6.0.2"
 from .base import Extractor
diff --git a/cognite/extractorutils/statestore.py b/cognite/extractorutils/statestore.py
index 41051c5c..14b959d8 100644
--- a/cognite/extractorutils/statestore.py
+++ b/cognite/extractorutils/statestore.py
@@ -92,10 +92,8 @@
 from types import TracebackType
 from typing import Any, Callable, Dict, Iterator, List, Optional, Tuple, Type, Union
 
-from requests.exceptions import ConnectionError
-
 from cognite.client import CogniteClient
-from cognite.client.exceptions import CogniteAPIError
+from cognite.client.exceptions import CogniteAPIError, CogniteException
 from cognite.extractorutils.uploader import DataPointList
 
 from ._inner_util import _DecimalDecoder, _DecimalEncoder, _resolve_log_level
@@ -348,7 +346,7 @@ def __init__(
         self._ensure_table()
 
     @retry(
-        exceptions=(CogniteAPIError, ConnectionError),
+        exceptions=(CogniteException,),
         tries=RETRIES,
         delay=RETRY_DELAY,
         max_delay=RETRY_MAX_DELAY,
@@ -370,7 +368,7 @@ def initialize(self, force: bool = False) -> None:
         self._initialize_implementation(force)
 
     @retry(
-        exceptions=(CogniteAPIError, ConnectionError),
+        exceptions=(CogniteException,),
         tries=RETRIES,
         delay=RETRY_DELAY,
         max_delay=RETRY_MAX_DELAY,
@@ -404,7 +402,7 @@ def synchronize(self) -> None:
         self._synchronize_implementation()
 
     @retry(
-        exceptions=(CogniteAPIError, ConnectionError),
+        exceptions=(CogniteException,),
         tries=RETRIES,
         delay=RETRY_DELAY,
         max_delay=RETRY_MAX_DELAY,
diff --git a/pyproject.toml b/pyproject.toml
index f63cf01c..fc5f3699 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "cognite-extractor-utils"
-version = "6.0.1"
+version = "6.0.2"
 description = "Utilities for easier development of extractors for CDF"
 authors = ["Mathias Lohne <mathias.lohne@cognite.com>"]
 license = "Apache-2.0"