diff --git a/darwin/torch/transforms.py b/darwin/torch/transforms.py
index 4308a50b0..458a4454c 100644
--- a/darwin/torch/transforms.py
+++ b/darwin/torch/transforms.py
@@ -332,7 +332,7 @@ def __call__(self, image, annotation: dict = None) -> tuple:
         transformed_data = self.transform(**albu_data)
         image, transformed_annotation = self._post_process(transformed_data, annotation)
 
-        if list(albu_data.keys()) == ["image"]:
+        if len(annotation) < 1:
             return image
 
         return image, transformed_annotation