diff --git a/tools/releasetools/add_img_to_target_files.py b/tools/releasetools/add_img_to_target_files.py
index ebaa867..44719ed 100755
--- a/tools/releasetools/add_img_to_target_files.py
+++ b/tools/releasetools/add_img_to_target_files.py
@@ -450,7 +450,7 @@
     if partition not in needed_partitions:
       continue
     assert (partition in common.AVB_PARTITIONS or
-            partition.startswith('vbmeta_')), \
+            partition in common.AVB_VBMETA_PARTITIONS), \
         'Unknown partition: {}'.format(partition)
     assert os.path.exists(path), \
         'Failed to find {} for {}'.format(path, partition)