diff --git a/arkindex_worker/worker/classification.py b/arkindex_worker/worker/classification.py
index d9dfaad1dadead0e1843172bb373e17ab2c8653c..6f17ee61eff4e70f4e5ab0b19c2e7d7978962284 100644
--- a/arkindex_worker/worker/classification.py
+++ b/arkindex_worker/worker/classification.py
@@ -2,6 +2,7 @@
 from apistar.exceptions import ErrorResponse
 
 from arkindex_worker import logger
+from arkindex_worker.cache import CachedElement
 from arkindex_worker.models import Element
 
 
@@ -60,8 +61,8 @@ class ClassificationMixin(object):
         Create a classification on the given element through API
         """
         assert element and isinstance(
-            element, Element
-        ), "element shouldn't be null and should be of type Element"
+            element, (Element, CachedElement)
+        ), "element shouldn't be null and should be an Element or CachedElement"
         assert ml_class and isinstance(
             ml_class, str
         ), "ml_class shouldn't be null and should be of type str"
diff --git a/tests/test_elements_worker/test_classifications.py b/tests/test_elements_worker/test_classifications.py
index 05e8d8d615a6cffb8de9ca3c3bd99548a14dcca1..9e1655aadefa438a4a781d57862d1d111e278bb6 100644
--- a/tests/test_elements_worker/test_classifications.py
+++ b/tests/test_elements_worker/test_classifications.py
@@ -159,7 +159,10 @@ def test_create_classification_wrong_element(mock_elements_worker):
             confidence=0.42,
             high_confidence=True,
         )
-    assert str(e.value) == "element shouldn't be null and should be of type Element"
+    assert (
+        str(e.value)
+        == "element shouldn't be null and should be an Element or CachedElement"
+    )
 
     with pytest.raises(AssertionError) as e:
         mock_elements_worker.create_classification(
@@ -168,7 +171,10 @@ def test_create_classification_wrong_element(mock_elements_worker):
             confidence=0.42,
             high_confidence=True,
         )
-    assert str(e.value) == "element shouldn't be null and should be of type Element"
+    assert (
+        str(e.value)
+        == "element shouldn't be null and should be an Element or CachedElement"
+    )
 
 
 def test_create_classification_wrong_ml_class(mock_elements_worker, responses):