diff --git a/core/modules/translation_entity/translation_entity.admin.inc b/core/modules/translation_entity/translation_entity.admin.inc
index b4213e4..59cd851 100644
--- a/core/modules/translation_entity/translation_entity.admin.inc
+++ b/core/modules/translation_entity/translation_entity.admin.inc
@@ -132,14 +132,11 @@ function translation_entity_translatable_batch($translatable, $field_name, &$con
         ->count()
         ->execute();
 
+      $context['sandbox']['max'] += $count;
       $context['sandbox']['progress_entity_type'][$entity_type] = 0;
       $context['sandbox']['max_entity_type'][$entity_type] = $count;
     }
 
-    // @todo Move this into the loop above as soon as entity query correctly
-    //   handles count queries on fields shared among different entities.
-    $context['sandbox']['max'] += $count;
-
     if ($context['sandbox']['max'] === 0) {
       // Nothing to do.
       $context['finished'] = 1;
@@ -152,7 +149,6 @@ function translation_entity_translatable_batch($translatable, $field_name, &$con
       continue;
     }
 
-    // Number of entities to be processed for each step.
     $info = entity_get_info($entity_type);
     $offset = $context['sandbox']['progress_entity_type'][$entity_type];
     $query = entity_query($entity_type);