diff --git a/gliner/data_processing/processor.py b/gliner/data_processing/processor.py
index cd8367d..27e856c 100644
--- a/gliner/data_processing/processor.py
+++ b/gliner/data_processing/processor.py
@@ -2022,7 +2022,7 @@ def create_relation_labels(
         if "span_label" in batch and batch["span_label"] is not None:
             batch_ents = (batch["span_label"] > 0).sum(-1)
         else:
-            batch_ents = span_mask.long().squeeze(-1).sum(-1)
+            batch_ents = span_mask.long().sum(-1)
 
         # Batch CPU transfer to avoid per-element .item() sync
         batch_ents_cpu = batch_ents.tolist()