diff --git a/examples/train_smolvla.sh b/examples/train_smolvla.sh
index fd093fe0..e3029494 100644
--- a/examples/train_smolvla.sh
+++ b/examples/train_smolvla.sh
@@ -14,7 +14,7 @@ export WORK=/home/mustafa_shukor
 # export DATA_DIR=$WORK/.cache/huggingface/datasets
 # export HF_LEROBOT_HOME=$WORK/.cache/huggingface/lerobot
 
-# export HF_DATASETS_OFFLINE=1 
+# export HF_DATASETS_OFFLINE=1
 # export HF_HUB_OFFLINE=1
 
 # export WANDB_CACHE_DIR=/lustre/fsn1/projects/rech/dyf/ugz83ue/wandb
@@ -54,5 +54,4 @@ rm -r $TRAIN_DIR
 CUDA_VISIBLE_DEVICES=2 python lerobot/scripts/train.py \
      --policy.type=$POLICY  \
      --dataset.repo_id=$REPO_ID \
-     --output_dir=$TRAIN_DIR 
-
+     --output_dir=$TRAIN_DIR
diff --git a/lerobot/common/policies/smolvla/modeling_smolvla.py b/lerobot/common/policies/smolvla/modeling_smolvla.py
index 418940fc..4c3a17c3 100644
--- a/lerobot/common/policies/smolvla/modeling_smolvla.py
+++ b/lerobot/common/policies/smolvla/modeling_smolvla.py
@@ -352,7 +352,7 @@ class SmolVLAPolicy(PreTrainedPolicy):
             )
         # Preprocess image features present in the batch
         for key in present_img_keys:
-            img = batch[key][:, i, :, :, :] if batch[key].ndim == 5 else batch[key]
+            img = batch[key][:, -1, :, :, :] if batch[key].ndim == 5 else batch[key]
             if self.config.resize_imgs_with_padding is not None:
                 img = resize_with_pad(img, *self.config.resize_imgs_with_padding, pad_value=0)