diff --git a/aria/model/modeling_aria.py b/aria/model/modeling_aria.py
index e737c60..4957fd9 100644
--- a/aria/model/modeling_aria.py
+++ b/aria/model/modeling_aria.py
@@ -172,12 +172,7 @@ def set_moe_z_loss_coeff(self, value):
         Args:
             value: The z-loss coefficient value to set.
         """
-        if hasattr(self.language_model, "set_z_loss_coeff"):
-            self.language_model.set_z_loss_coeff(value)
-        else:
-            logger.warning(
-                "The language model does not have a `set_z_loss_coeff` method. Ignore this warning if you are not using a MoYI model."
-            )
+        self.language_model.set_z_loss_coeff(value)
 
     def set_moe_aux_loss_coeff(self, value):
         """
@@ -186,12 +181,7 @@ def set_moe_aux_loss_coeff(self, value):
         Args:
             value: The auxiliary loss coefficient value to set.
         """
-        if hasattr(self.language_model, "set_aux_loss_coeff"):
-            self.language_model.set_aux_loss_coeff(value)
-        else:
-            logger.warning(
-                "The language model does not have a `set_aux_loss_coeff` method. Ignore this warning if you are not using a MoYI model."
-            )
+        self.language_model.set_aux_loss_coeff(value)
 
     # copied from transformers.models.llava.modeling_llava.LlavaForConditionalGeneration
     def _merge_input_ids_with_image_features(
diff --git a/aria/model/vision_processor.py b/aria/model/vision_processor.py
index fe263dd..c9e60bd 100644
--- a/aria/model/vision_processor.py
+++ b/aria/model/vision_processor.py
@@ -98,6 +98,11 @@ def __init__(
         self.min_image_size = min_image_size
         self.image_mean = image_mean
         self.image_std = image_std
+        self.auto_map = {
+            "AutoProcessor": "processing_aria.AriaProcessor",
+            "AutoImageProcessor": "vision_processor.AriaVisionProcessor",
+        }
+
         # we make the transform a property so that it is lazily initialized,
         # this could avoid the error "TypeError: Object of type Normalize is not JSON serializable"
         # when we used save_pretrained or from_pretrained.