diff --git a/src/maxtext/trainers/post_train/sft/train_sft.py b/src/maxtext/trainers/post_train/sft/train_sft.py
index 9002caa16c..e1cf97c62d 100644
--- a/src/maxtext/trainers/post_train/sft/train_sft.py
+++ b/src/maxtext/trainers/post_train/sft/train_sft.py
@@ -210,6 +210,9 @@ def loss_func(
       targets,
       targets_position,
       targets_segmentation,
+      images=None,
+      image_masks=None,
+      **kwargs,
   ):
     data = {
         "inputs": inputs,
@@ -219,6 +222,10 @@ def loss_func(
         "targets_position": targets_position,
         "targets_segmentation": targets_segmentation,
     }
+    if images is not None:
+      data["images"] = images
+    if image_masks is not None:
+      data["image_masks"] = image_masks
     return loss_fn(model, mt_config, data, dropout_rng=None, params=None, is_train=True)
 
   trainer = trainer.with_loss_fn(loss_func, has_aux=True)