diff --git a/.gitignore b/.gitignore
index 4e22c95..6b921ea 100644
--- a/.gitignore
+++ b/.gitignore
@@ -2,7 +2,6 @@ model/clip_model/clip-vit-base-patch32/
 out/*.pth
 full.json
 trans_json.py
-dataset/sft_multi_images/
-dataset/
-!dataset/eval_images
-!dataset/eval_multi_images
\ No newline at end of file
+dataset/*
+!dataset/eval_images/
+!dataset/eval_multi_images/
\ No newline at end of file
diff --git a/2-sft_vlm.py b/2-sft_vlm.py
index 321dfd6..4bc5210 100644
--- a/2-sft_vlm.py
+++ b/2-sft_vlm.py
@@ -18,7 +18,7 @@ from torch.utils.data import DataLoader, DistributedSampler
 from transformers import AutoTokenizer, AutoModel
 from model.model import Transformer
 from model.LMConfig import LMConfig
-from model.dataset import SFTDataset
+from model.dataset import SFTDataset, SFTDataset_multi
 from model.vision_utils import get_vision_model, get_img_embedding
 
 warnings.filterwarnings('ignore')
@@ -97,7 +97,10 @@ def train_epoch(epoch, wandb):
         if (step + 1) % args.save_interval == 0 and (not ddp or dist.get_rank() == 0):
             model.eval()
             moe_path = '_moe' if lm_config.use_moe else ''
-            ckp = f'{args.save_dir}/{lm_config.dim}{moe_path}_vlm_sft.pth'
+            if args.multi:
+                ckp = f'{args.save_dir}/{lm_config.dim}{moe_path}_vlm_sft_multi.pth'
+            else:
+                ckp = f'{args.save_dir}/{lm_config.dim}{moe_path}_vlm_sft.pth'
 
             if isinstance(model, torch.nn.parallel.DistributedDataParallel):
                 state_dict = model.module.state_dict()
@@ -106,12 +109,29 @@ def train_epoch(epoch, wandb):
 
             torch.save(state_dict, ckp)
             model.train()
+    if args.save_last and (not ddp or dist.get_rank() == 0) and args.multi:
+        model.eval()
+        moe_path = '_moe' if lm_config.use_moe else ''
+        ckp = f'{args.save_dir}/{lm_config.dim}{moe_path}_vlm_sft_multi.pth'
+
+        if isinstance(model, torch.nn.parallel.DistributedDataParallel):
+            state_dict = model.module.state_dict()
+        else:
+            state_dict = model.state_dict()
+
+        torch.save(state_dict, ckp)
+        model.train()
 
 
 def init_model(lm_config):
     tokenizer = AutoTokenizer.from_pretrained('./model/minimind_tokenizer')
     moe_path = '_moe' if lm_config.use_moe else ''
-    ckp = f'./out/{lm_config.dim}{moe_path}_vlm_pretrain.pth'
+
+    # 多图推理建议在单图推理之后
+    if args.multi:
+        ckp = f'./out/{lm_config.dim}{moe_path}_vlm_sft.pth'
+    else:
+        ckp = f'./out/{lm_config.dim}{moe_path}_vlm_pretrain.pth'
 
     model = Transformer(lm_config)
     state_dict = torch.load(ckp, map_location=args.device)
@@ -159,6 +179,8 @@ if __name__ == "__main__":
     parser.add_argument("--num_workers", type=int, default=8, help="Number of workers for data loading")
     parser.add_argument("--data_path", type=str, default="./dataset/LLaVA-Instruct/llava_instruct_230k.json",
                         help="Path to training data")
+    parser.add_argument("--data_path_multi", type=str, default="./dataset/sft_multi_images/output.json",
+                        help="Path to multi images training data")
     parser.add_argument("--ddp", action="store_true", help="Use DistributedDataParallel")
     parser.add_argument("--accumulation_steps", type=int, default=1, help="Gradient accumulation steps")
     parser.add_argument("--grad_clip", type=float, default=1.0, help="Gradient clipping threshold")
@@ -166,6 +188,8 @@ if __name__ == "__main__":
     parser.add_argument("--log_interval", type=int, default=100, help="Logging interval")
     parser.add_argument("--save_interval", type=int, default=100, help="Model saving interval")
     parser.add_argument('--local_rank', type=int, default=-1, help='local rank for distributed training')
+    parser.add_argument('--multi', type=bool, default=False, help='multi-images training')
+    parser.add_argument('--save_last', type=bool, default=True, help='save last step model')
 
     args = parser.parse_args()
 
@@ -197,9 +221,15 @@ if __name__ == "__main__":
     model, tokenizer, (vision_model, preprocess) = init_model(lm_config)
 
     use_version = 0
-    train_ds = SFTDataset(args.data_path, tokenizer, vision_model=(vision_model, preprocess),
+    if args.multi:
+        print("进行多图训练，建议在指令微调后进行...")
+        train_ds = SFTDataset_multi(args.data_path_multi, tokenizer, vision_model=(vision_model, preprocess),
                           image_special_token=lm_config.image_special_token,
                           max_length=max_seq_len)
+    else:
+        train_ds = SFTDataset(args.data_path, tokenizer, vision_model=(vision_model, preprocess),
+                            image_special_token=lm_config.image_special_token,
+                            max_length=max_seq_len)
     train_sampler = DistributedSampler(train_ds) if ddp else None
     train_loader = DataLoader(
         train_ds,
diff --git a/3-eval_chat.py b/3-eval_chat.py
index b0d8301..7f1f214 100644
--- a/3-eval_chat.py
+++ b/3-eval_chat.py
@@ -1,11 +1,8 @@
 import os
 import random
-import time
-
 import numpy as np
 import torch
 import warnings
-
 from PIL import Image
 from transformers import AutoTokenizer, AutoModelForCausalLM
 from model.model import Transformer
@@ -19,14 +16,16 @@ def count_parameters(model):
     return sum(p.numel() for p in model.parameters() if p.requires_grad)
 
 
-def init_model(lm_config):
+def init_model(lm_config, device, multi):
     tokenizer = AutoTokenizer.from_pretrained('./model/minimind_tokenizer')
     model_from = 1  # 1从权重，2用transformers
 
     if model_from == 1:
         moe_path = '_moe' if lm_config.use_moe else ''
-        ckp = f'./out/{lm_config.dim}{moe_path}_vlm_sft.pth'
-
+        if multi:
+            ckp = f'./out/{lm_config.dim}{moe_path}_vlm_sft_multi.pth'
+        else:
+            ckp = f'./out/{lm_config.dim}{moe_path}_vlm_sft.pth'
         model = Transformer(lm_config)
         state_dict = torch.load(ckp, map_location=device)
 
@@ -35,125 +34,103 @@ def init_model(lm_config):
         for k, v in list(state_dict.items()):
             if k.startswith(unwanted_prefix):
                 state_dict[k[len(unwanted_prefix):]] = state_dict.pop(k)
-
-        # 加载到模型中
         model.load_state_dict(state_dict, strict=False)
     else:
-        model = AutoModelForCausalLM.from_pretrained('minimind-v-v1-small',
-                                                     trust_remote_code=True)
+        model = AutoModelForCausalLM.from_pretrained('minimind-v-v1-small', trust_remote_code=True)
 
     model = model.to(device)
-
     print(f'模型参数: {count_parameters(model) / 1e6} 百万 = {count_parameters(model) / 1e9} B (Billion)')
 
-    (vision_model, preprocess) = get_vision_model()
+    vision_model, preprocess = get_vision_model()
     vision_model = vision_model.to(device)
-    return model, tokenizer, (vision_model, preprocess)
+    return model, tokenizer, vision_model, preprocess
 
 
 def setup_seed(seed):
-    random.seed(seed)  # 设置 Python 的随机种子
-    np.random.seed(seed)  # 设置 NumPy 的随机种子
-    torch.manual_seed(seed)  # 设置 PyTorch 的随机种子
-    torch.cuda.manual_seed(seed)  # 为当前 GPU 设置随机种子（如果有）
-    torch.cuda.manual_seed_all(seed)  # 为所有 GPU 设置随机种子（如果有）
-    torch.backends.cudnn.deterministic = True  # 确保每次返回的卷积算法是确定的
-    torch.backends.cudnn.benchmark = False  # 关闭 cuDNN 的自动调优，避免不确定性
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)
+    torch.backends.cudnn.deterministic = True
+    torch.backends.cudnn.benchmark = False
 
 
 if __name__ == "__main__":
-    # -----------------------------------------------------------------------------
+    # -------------------------- 基本参数设置 -----------------------------------
+    multi = False  # 设置multi参数，控制单图或多图推理
     out_dir = 'out'
-    start = ""
     temperature = 0.5
     top_k = 8
     setup_seed(1337)
-    # device = 'cpu'
     device = 'cuda:0' if torch.cuda.is_available() else 'cpu'
     dtype = 'bfloat16'
-    max_seq_len = 1 * 1024
+    max_seq_len = 1024
     lm_config = LMConfig()
     lm_config.max_seq_len = max_seq_len
-    # -----------------------------------------------------------------------------
+    model, tokenizer, vision_model, preprocess = init_model(lm_config, device, multi)
+    model.eval()
+    
+    # -------------------------- 问题和目录设置 -----------------------------------
+    if multi:
+        image_dir = './dataset/eval_multi_images/bird/'
+        prompt = "<image>\n<image>\nName all the differences between these two birds."
+    else:
+        image_dir = './dataset/eval_images/'
+        prompt = lm_config.image_special_token + '\n这个图片描述的是什么内容？'
 
-    model, tokenizer, (vision_model, preprocess) = init_model(lm_config)
-
-    model = model.eval()
-    # 推送到huggingface
-    # model.push_to_hub("minimind")
-    # tokenizer.push_to_hub("minimind")
-
-    stream = True
-
-    prompt = lm_config.image_special_token + '\n这个图片描述的是什么内容？'
-
-    i = 0
-    # 获取图像文件列表
-    image_dir = './dataset/eval_images/'
     image_files = sorted(os.listdir(image_dir))
 
-    # 遍历图像文件
-    for image_file in image_files:
-        image_path = os.path.join(image_dir, image_file)
-
-        # 打开图像并转换为RGB格式
-        image = Image.open(image_path).convert('RGB')
-        image_process = get_img_process(image, preprocess).to(vision_model.device)
-        # 对图像进行编码
-        image_encoder = get_img_embedding(image_process, vision_model).unsqueeze(0)
-
-        print(f'[Image]: {image_file}')
-        prompt_ = prompt.replace("\n", "")
-        print(f'[Q]: {prompt_}')
+    # -------------------------- 推理逻辑 -----------------------------------
+    if multi:
+        # 多图推理：所有图像编码一次性推理
+        image_encoders = []
+        for image_file in image_files:
+            image_path = os.path.join(image_dir, image_file)
+            image = Image.open(image_path).convert('RGB')
+            image_process = get_img_process(image, preprocess).to(vision_model.device)
+            image_encoder = get_img_embedding(image_process, vision_model).unsqueeze(0)
+            image_encoders.append(image_encoder)
+            print(f'[Image]: {image_file}')
+        image_encoders = torch.cat(image_encoders, dim=0).unsqueeze(0)
 
         messages = [{"role": "user", "content": prompt}]
-
-        # print(messages)
-        new_prompt = tokenizer.apply_chat_template(
-            messages,
-            tokenize=False,
-            add_generation_prompt=True
-        )[-(max_seq_len - 1):]
-
+        new_prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)[-(max_seq_len - 1):]
         x = tokenizer(new_prompt).data['input_ids']
-        x = (torch.tensor(x, dtype=torch.long, device=device)[None, ...])
-
-        answer = new_prompt
+        x = torch.tensor(x, dtype=torch.long, device=device)[None, ...]
 
         with torch.no_grad():
             res_y = model.generate(x, tokenizer.eos_token_id, max_new_tokens=max_seq_len, temperature=temperature,
-                                   top_k=top_k, stream=stream, image_encoders=image_encoder)
+                                   top_k=top_k, stream=True, image_encoders=image_encoders)
             print('[A]: ', end='')
-            try:
-                y = next(res_y)
-            except StopIteration:
-                print("No answer")
-                continue
-
             history_idx = 0
-            while y != None:
+            for y in res_y:
                 answer = tokenizer.decode(y[0].tolist())
-                if answer and answer[-1] == '�':
-                    try:
-                        y = next(res_y)
-                    except:
-                        break
-                    continue
-                # print(answer)
-                if not len(answer):
-                    try:
-                        y = next(res_y)
-                    except:
-                        break
-                    continue
-
                 print(answer[history_idx:], end='', flush=True)
-                try:
-                    y = next(res_y)
-                except:
-                    break
                 history_idx = len(answer)
-                if not stream:
-                    break
-
             print('\n')
+
+    else:
+        # 单图推理：对每张图像单独推理
+        for image_file in image_files:
+            image_path = os.path.join(image_dir, image_file)
+            image = Image.open(image_path).convert('RGB')
+            image_process = get_img_process(image, preprocess).to(vision_model.device)
+            image_encoder = get_img_embedding(image_process, vision_model).unsqueeze(0)
+
+            messages = [{"role": "user", "content": prompt}]
+            new_prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)[-(max_seq_len - 1):]
+            x = tokenizer(new_prompt).data['input_ids']
+            x = torch.tensor(x, dtype=torch.long, device=device)[None, ...]
+
+            print(f'[Image]: {image_file}')
+            with torch.no_grad():
+                res_y = model.generate(x, tokenizer.eos_token_id, max_new_tokens=max_seq_len, temperature=temperature,
+                                       top_k=top_k, stream=True, image_encoders=image_encoder)
+                print('[A]: ', end='')
+                history_idx = 0
+                for y in res_y:
+                    answer = tokenizer.decode(y[0].tolist())
+                    print(answer[history_idx:], end='', flush=True)
+                    history_idx = len(answer)
+                print('\n')
\ No newline at end of file
diff --git a/5-eval_chat_multi.py b/5-eval_chat_multi.py
index 9134e10..59b4c5e 100644
--- a/5-eval_chat_multi.py
+++ b/5-eval_chat_multi.py
@@ -85,8 +85,6 @@ if __name__ == "__main__":
 
     stream = True
 
-    # prompt = lm_config.image_special_token + '\n这个图片描述的是什么内容？'
-
     i = 0
     # 获取图像文件列表
     # image_dir = './dataset/eval_multi_images/' # 对于图像转换数据集
diff --git a/README.md b/README.md
index ca42ed8..1883d74 100644
--- a/README.md
+++ b/README.md
@@ -180,7 +180,7 @@ CPU: Intel(R) Core(TM) i9-10980XE CPU @ 3.00GHz
       目录下，命名为`*_llm.pth`
     * 3.4 `python 1-pretrain_vlm.py` 执行预训练，得到 `*_vlm_pretrain.pth` 作为预训练的输出权重
     * 3.5 `python 2-sft_vlm.py` 执行指令微调，得到 `*_vlm_sft.pth` 作为指令微调的输出权重
-    * 3.6 `python 4-sft_multi_vlm.py` 执行多图指令微调（效果仅参考），得到 `*_vlm_sft_multi.pth` 作为多图指令微调的输出权重，512+8 模型显存占用约为8198M
+    * 3.6 `python 2-sft_vlm.py --multi True` 在指令微调的基础上执行多图指令微调（效果仅参考），得到 `*_vlm_sft_multi.pth` 作为多图指令微调的输出权重，512+8 模型显存占用约为8198M
 
 * 4.测试自己训练的模型推理效果
     * 确保需要使用的，训练完成的参数权重`*.pth`文件位于`./out/`目录下
@@ -196,10 +196,7 @@ CPU: Intel(R) Core(TM) i9-10980XE CPU @ 3.00GHz
       ```
     * `python 3-eval_chat.py`测试模型的对话效果，其中测试图片在`./dataset/eval_images`下，可自行更换
       ![eval_chat](images/3-eval_chat.png)
-<<<<<<< HEAD
-    * `python 5-eval_chat_multi.py`测试模型的多图对话效果，其中测试图片在`./dataset/eval_multi_images`下，可自行更换（多图数据集规模相对较小且为英文对话，数据集仅包含两图对比的场景，因此微调效果有限）。
-=======
->>>>>>> 606d58814fb72b49674ce4fdd47aed1e7f36d827
+    * `python 3-eval_chat.py`调整[multi](./3-eval_chat.py#L61)变量，可以测试模型的多图对话效果，其中测试图片在`./dataset/eval_multi_images`下，可自行更换（多图数据集规模相对较小且为英文对话，数据集仅包含两图对比的场景，因此微调效果有限）。
 
 🍭 【Tip】预训练和全参指令微调pretrain和sft均支持多卡加速
 
@@ -316,7 +313,6 @@ minimind-v使用50个字符组成的 `<<<...>>>` 占位符代替图像，
 因此minimind-v的prompt：
 
 ```text
-<<<<<<<<<<<<<<<<<<<<<<<<<>>>>>>>>>>>>>>>>>>>>>>>>>\n这个图片描述的是什么内容？
 ```
 
 计算完embedding和projection，并对图像部分token替换后
@@ -324,14 +320,12 @@ minimind-v使用50个字符组成的 `<<<...>>>` 占位符代替图像，
 
 ![input](./images/minimind-v-input.png)
 
-<<<<<<< HEAD
+
 多图实现方法就是通过注入多个\<image\>图像占位符进行实现，不需要修改任何框架。
 
->  ps: 唯一值得注意的点是，如果在训练过程中存在不同conversations插入图片数量不同的情况，需要利用空特征将较短的特征进行填充（对应[dataset的第267行](./model/dataset.py)），以保证能够在同样大小下被dataloader读取。
+>  ps: 唯一值得注意的点是，如果在训练过程中存在不同conversations插入图片数量不同的情况，需要利用空特征将较短的特征进行填充（对应[dataset的第267行](./model/dataset.py#267)），以保证能够在同样大小下被dataloader读取。
 > pps: 在prompt中不需要如此做，仍旧是根据插入图像的数量来进行占位符的注入。因此，最终输入给LLM的input feature不会受填充特征的影响。
 
-=======
->>>>>>> 606d58814fb72b49674ce4fdd47aed1e7f36d827
 <u>至此，MiniMind-V的所有细节已经呈现完毕。</u>
 
 <u>MiniMind-V的实现未参考 **任何** 第三方代码，完全基于MiniMind尽可能做最小改动产生，故代码实现和LlaVA等模型必然存在很大区别。
@@ -382,7 +376,6 @@ MiniMind-V与MiniMind的代码核心改动不超过100行，上手难度低。</
 }
 ```
 
-<<<<<<< HEAD
 多图指令微调数据集格式：
 
 ```json
@@ -422,32 +415,23 @@ MiniMind-V与MiniMind的代码核心改动不超过100行，上手难度低。</
 注：
 + 对于指令微调，仅保留了一轮对话，训练单轮对话模型，防止小模型性能被长文本拉低。
 + 多图数据集规模相对较小且为英文对话，数据集仅包含两图对比的场景，因此微调效果有限，这里只提供一种参考思路。
-=======
-注：对于指令微调，仅保留了一轮对话，训练单轮对话模型，防止小模型性能被长文本拉低。
->>>>>>> 606d58814fb72b49674ce4fdd47aed1e7f36d827
 
 最终的数据集下载地址：[百度网盘](https://pan.baidu.com/s/1Nz36OBBvVBGEx-PwIb7ofg?pwd=6666) | [HuggingFace](https://huggingface.co/datasets/jingyaogong/minimind-v_dataset)
 
+多图对话数据集：[HuggingFace](https://hf-mirror.com/datasets/xinyanghuang/minimind-v_multi_image/tree/main)
+
 ## 训练
 
 预训练从595K条数据集中学习图片的通用知识，比如鹿是鹿，狗是狗。
 
 指令微调从230K条真实对话数据集中学习对图片提问的真实问答格式。
 
-<<<<<<< HEAD
-多图微调提供两个数据集，长度分别为3.5k和13.6k的真实问答格式。
+多图微调提供两个数据集，图像转换数据集和鸟类对比数据集，长度分别为3.5k和13.6k的真实问答格式。
 
-=======
->>>>>>> 606d58814fb72b49674ce4fdd47aed1e7f36d827
 `1-pretrain_vlm.py` 执行预训练，得到 `*_vlm_pretrain.pth` 作为预训练的输出权重。
 
 `2-sft_vlm.py` 执行指令微调，得到 `*_vlm_sft.pth` 作为指令微调的输出权重。
 
-<<<<<<< HEAD
-`4-sft_multi_vlm.py` 执行指令微调，得到 `*_vlm_sft_multi.pth` 作为指令微调的输出权重。
-
-=======
->>>>>>> 606d58814fb72b49674ce4fdd47aed1e7f36d827
 训练时均冻结visual encoder也就是clip模型，只微调Projection和LLM两部分。
 
 > Pretrain 512+8 模型 (训练时间和Loss参考图)
@@ -477,10 +461,7 @@ MiniMind-V与MiniMind的代码核心改动不超过100行，上手难度低。</
 
 ### 效果测试
 
-<<<<<<< HEAD
 #### 单图对话
-=======
->>>>>>> 606d58814fb72b49674ce4fdd47aed1e7f36d827
 <table>
   <thead>
     <tr>
@@ -565,7 +546,6 @@ MiniMind-V与MiniMind的代码核心改动不超过100行，上手难度低。</
   </tbody>
 </table>
 
-<<<<<<< HEAD
 #### 多图对话（效果十分有限）
 
 <table>
@@ -578,15 +558,13 @@ MiniMind-V与MiniMind的代码核心改动不超过100行，上手难度低。</
   </thead>
   <tbody>
     <tr>
-      <td><img src="./dataset/eval_multi_images/bird//0.jpg" alt="a-bird.png" style="width: 200px;"></td>
-      <td><img src="./dataset/eval_multi_images/bird//1.jpg" alt="a-bird.png" style="width: 200px;"></td>
+      <td><img src="./dataset/eval_multi_images/bird/0.jpg" alt="a-bird.png" style="width: 200px;"></td>
+      <td><img src="./dataset/eval_multi_images/bird/1.jpg" alt="a-bird.png" style="width: 200px;"></td>
       <td>animal1 has a brown and black head with a black and white striped head . animal2 has a black head with a white stripe on its wings .</td>
     </tr>
   </tbody>
 </table>
 
-=======
->>>>>>> 606d58814fb72b49674ce4fdd47aed1e7f36d827
 ### 启动推理
 
 ```bash
diff --git a/dataset/eval_multi_images/bird/0.jpg b/dataset/eval_multi_images/bird/0.jpg
new file mode 100644
index 0000000..fa2aa25
--- /dev/null
+++ b/dataset/eval_multi_images/bird/0.jpg
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d34d90f79ad34a62ed46bf2ddcd2b348703c04817befceb6ece6629f508b8636
+size 57989
diff --git a/dataset/eval_multi_images/bird/1.jpg b/dataset/eval_multi_images/bird/1.jpg
new file mode 100644
index 0000000..7628aab
--- /dev/null
+++ b/dataset/eval_multi_images/bird/1.jpg
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:86f21fb466259b028c536a9814aee11165348df2f7443dff4e069b53b6c05a5f
+size 96613
diff --git a/model/__pycache__/LMConfig.cpython-310.pyc b/model/__pycache__/LMConfig.cpython-310.pyc
index 8f08b70..3b4c852 100644
Binary files a/model/__pycache__/LMConfig.cpython-310.pyc and b/model/__pycache__/LMConfig.cpython-310.pyc differ
diff --git a/model/__pycache__/dataset.cpython-310.pyc b/model/__pycache__/dataset.cpython-310.pyc
index ea280d7..ef0806e 100644
Binary files a/model/__pycache__/dataset.cpython-310.pyc and b/model/__pycache__/dataset.cpython-310.pyc differ
diff --git a/model/__pycache__/model.cpython-310.pyc b/model/__pycache__/model.cpython-310.pyc
index 2ed9f3e..2df8cd0 100644
Binary files a/model/__pycache__/model.cpython-310.pyc and b/model/__pycache__/model.cpython-310.pyc differ
diff --git a/model/__pycache__/vision_utils.cpython-310.pyc b/model/__pycache__/vision_utils.cpython-310.pyc
index 0bf015b..97b6ac5 100644
Binary files a/model/__pycache__/vision_utils.cpython-310.pyc and b/model/__pycache__/vision_utils.cpython-310.pyc differ
diff --git a/model/model.py b/model/model.py
index 35b2f00..f985aaa 100644
--- a/model/model.py
+++ b/model/model.py
@@ -385,29 +385,26 @@ class Transformer(PreTrainedModel):
 
         # 如果此时有图像编码
         if image_encoders is not None:
-            vision_proj = self.vision_proj(image_encoders)  # torch.Size([32, 4, 50, 512])
-            if image_indices is not None:
-                # 创建一个新的张量来存储拼接后的结果
-                new_h = []
-
-                for i in range(h.size(0)):
-                    current_h = h[i]
-                    img_idx = 0
-                    for _, (batch_idx, start_idx, end_idx) in enumerate(image_indices):
-                        if batch_idx == i:
-                            # 插入vision_proj特征
-                            before = current_h[:start_idx, :]
-                            after = current_h[end_idx + 1:, :]
-                            # 拼接 before, vision_proj, after
-                            if len(vision_proj.shape) == 4:
-                                current_h = torch.cat((before, vision_proj[i][img_idx], after), dim=0)[:seqlen]
-                            else:
-                                current_h = torch.cat((before, vision_proj[i], after), dim=0)[:seqlen]
-                            img_idx += 1
-                    new_h.append(current_h)
-                # 将所有拼接后的结果堆叠起来
-                new_h = torch.stack(new_h, dim=0) # torch.Size([32, 511, 512])
-                return new_h
+                vision_proj = self.vision_proj(image_encoders)
+                vision_proj = vision_proj.unsqueeze(0) if len(vision_proj.shape) == 3 else vision_proj
+                if image_indices is not None:
+                    # 创建一个新的张量来存储拼接后的结果
+                    new_h = []
+                    for i in range(h.size(0)):
+                        # i即为current_batch_idx索引
+                        img_idx = 0
+                        for batch_idx, start_idx, end_idx in image_indices:
+                            if batch_idx == i:
+                                # 插入vision_proj特征
+                                before = h[i][:start_idx, :]
+                                after = h[i][end_idx + 1:, :]
+                                # 拼接 before, vision_proj, after
+                                h[i] = torch.cat((before, vision_proj[i][img_idx], after), dim=0)[:seqlen]
+                                img_idx += 1
+                        new_h.append(h[i])
+                    # 将所有拼接后的结果重新堆叠起来
+                    new_h = torch.stack(new_h, dim=0)  # torch.Size([32, 511, 512])
+                    return new_h
 
         return h