IndexError: The shape of the mask [6, 440] at index 1 does not match the shape of the indexed tensor [6, 760, 256] at index 1

[Oseemaker]

Hello, I use the code deepspeed --master_port=24999 train_ds.py --version="/dat02/sh/LLaVA-Med-1.0.0/models/llava_med_model" --dataset_dir='./dataset' --vision_pretrained="/dat02/sh/LISA-main/weights/sam_vit_h_4b8939.pth" - the dataset = "refer_seg | | reason_seg" - sample_rates = "3, 1" - exp_name = "lisa - 7 b"
Wanted to retrain a new model on llava, but encountered a new problem:
File "/dat02/sh/LISA-main/train_ds.py", line 589, in
main(sys.argv[1:])
File "/dat02/sh/LISA-main/train_ds.py", line 374, in main
train_iter = train(
File "/dat02/sh/LISA-main/train_ds.py", line 463, in train
output_dict = model(**input_dict)
File "/dat01/anaconda3/envs/lisa2/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1194, in _call_impl
return forward_call(*input, **kwargs)
File "/dat01/anaconda3/envs/lisa2/lib/python3.10/site-packages/deepspeed/utils/nvtx.py", line 15, in wrapped_fn
ret_val = func(*args, **kwargs)
File "/dat01/anaconda3/envs/lisa2/lib/python3.10/site-packages/deepspeed/runtime/engine.py", line 1846, in forward
loss = self.module(*inputs, **kwargs)
File "/dat01/anaconda3/envs/lisa2/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1194, in _call_impl
return forward_call(*input, **kwargs)
File "/dat01/anaconda3/envs/lisa2/lib/python3.10/site-packages/peft/peft_model.py", line 922, in forward
return self.base_model(
File "/dat01/anaconda3/envs/lisa2/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1194, in _call_impl
return forward_call(*input, **kwargs)
File "/dat02/sh/LISA-main/model/LISA.py", line 167, in forward
return self.model_forward(**kwargs)
File "/dat02/sh/LISA-main/model/LISA.py", line 256, in model_forward
pred_embeddings = last_hidden_state[seg_token_mask]
IndexError: The shape of the mask [6, 440] at index 1 does not match the shape of the indexed tensor [6, 760, 256] at index 1