diff --git a/tools/checkpoint_conversion/convert_vit_checkpoints.py b/tools/checkpoint_conversion/convert_vit_checkpoints.py
index cd382b2165..cb6f43702b 100644
--- a/tools/checkpoint_conversion/convert_vit_checkpoints.py
+++ b/tools/checkpoint_conversion/convert_vit_checkpoints.py
@@ -4,7 +4,7 @@
 export KAGGLE_KEY=XXX
 
 python tools/checkpoint_conversion/convert_vit_checkpoints.py \
-    --preset vit_base_patch16_224
+    --preset vit_base_patch16_224_imagenet
 """
 
 import os
@@ -30,17 +30,17 @@
 FLAGS = flags.FLAGS
 
 PRESET_MAP = {
-    "vit_base_patch16_224": "google/vit-base-patch16-224",
-    "vit_base_patch16_384": "google/vit-base-patch16-384",
-    "vit_base_patch32_384": "google/vit-base-patch32-384",
-    "vit_large_patch16_224": "google/vit-large-patch16-224",
-    "vit_large_patch16_384": "google/vit-large-patch16-384",
-    "vit_large_patch32_384": "google/vit-large-patch32-384",
-    "vit_base_patch16_224_in21k": "google/vit-base-patch16-224-in21k",
-    "vit_base_patch32_224_in21k": "google/vit-base-patch32-224-in21k",
-    "vit_large_patch16_224_in21k": "google/vit-large-patch16-224-in21k",
-    "vit_large_patch32_224_in21k": "google/vit-large-patch32-224-in21k",
-    "vit_huge_patch14_224_in21k": "google/vit-huge-patch14-224-in21k",
+    "vit_base_patch16_224_imagenet": "google/vit-base-patch16-224",
+    "vit_base_patch16_384_imagenet": "google/vit-base-patch16-384",
+    "vit_base_patch32_384_imagenet": "google/vit-base-patch32-384",
+    "vit_large_patch16_224_imagenet": "google/vit-large-patch16-224",
+    "vit_large_patch16_384_imagenet": "google/vit-large-patch16-384",
+    "vit_large_patch32_384_imagenet": "google/vit-large-patch32-384",
+    "vit_base_patch16_224_imagenet21k": "google/vit-base-patch16-224-in21k",
+    "vit_base_patch32_224_imagenet21k": "google/vit-base-patch32-224-in21k",
+    "vit_large_patch16_224_imagenet21k": "google/vit-large-patch16-224-in21k",
+    "vit_large_patch32_224_imagenet21k": "google/vit-large-patch32-224-in21k",
+    "vit_huge_patch14_224_imagenet21k": "google/vit-huge-patch14-224-in21k",
 }
 
 flags.DEFINE_string(