diff --git a/tools/releasetools/add_img_to_target_files.py b/tools/releasetools/add_img_to_target_files.py
index e75b3b7..f545dcc 100755
--- a/tools/releasetools/add_img_to_target_files.py
+++ b/tools/releasetools/add_img_to_target_files.py
@@ -119,9 +119,12 @@
 
   simg = sparse_img.SparseImage(imgname)
   care_map_ranges = simg.care_map
-  key = which + "_image_blocks"
-  image_blocks = OPTIONS.info_dict.get(key)
-  if image_blocks:
+  size_key = which + "_image_size"
+  image_size = OPTIONS.info_dict.get(size_key)
+  if image_size:
+    # excludes the verity metadata blocks of the given image. When AVB is enabled,
+    # this size is the max image size returned by the AVB tool
+    image_blocks = int(image_size) / 4096 - 1
     assert image_blocks > 0, "blocks for {} must be positive".format(which)
     care_map_ranges = care_map_ranges.intersect(
         rangelib.RangeSet("0-{}".format(image_blocks)))
@@ -319,9 +322,7 @@
   if block_list:
     block_list.Write()
 
-  # Set the '_image_blocks' that excludes the verity metadata blocks of the
-  # given image. When AVB is enabled, this size is the max image size returned
-  # by the AVB tool.
+  # Set the '_image_size' for given image size.
   is_verity_partition = "verity_block_device" in image_props
   verity_supported = (image_props.get("verity") == "true" or
                       image_props.get("avb_enable") == "true")
@@ -329,8 +330,8 @@
   if verity_supported and (is_verity_partition or is_avb_enable):
     image_size = image_props.get("image_size")
     if image_size:
-      image_blocks_key = what + "_image_blocks"
-      info_dict[image_blocks_key] = int(image_size) / 4096 - 1
+      image_size_key = what + "_image_size"
+      info_dict[image_size_key] = int(image_size)
 
   use_dynamic_size = (
       info_dict.get("use_dynamic_partition_size") == "true" and
diff --git a/tools/releasetools/test_add_img_to_target_files.py b/tools/releasetools/test_add_img_to_target_files.py
index d2a274d..482f86c 100644
--- a/tools/releasetools/test_add_img_to_target_files.py
+++ b/tools/releasetools/test_add_img_to_target_files.py
@@ -362,7 +362,7 @@
         (0xCAC3, 4),
         (0xCAC1, 6)])
     OPTIONS.info_dict = {
-        'system_image_blocks' : 12,
+        'system_image_size' : 53248,
     }
     name, care_map = GetCareMap('system', sparse_image)
     self.assertEqual('system', name)
@@ -377,6 +377,6 @@
         (0xCAC3, 4),
         (0xCAC1, 6)])
     OPTIONS.info_dict = {
-        'system_image_blocks' : -12,
+        'system_image_size' : -45056,
     }
     self.assertRaises(AssertionError, GetCareMap, 'system', sparse_image)