Fix tensor input handling in MiniMindVLM.get_image_embeddings

pull/130/head
ian 2026-09-30 18:10:37 +08:00
parent 1862b633fc
commit e2332de909
1 changed files with 5 additions and 3 deletions

View File

@ -66,10 +66,12 @@ class MiniMindVLM(MiniMindForCausalLM):
@staticmethod @staticmethod
def get_image_embeddings(image_inputs, vision_model): def get_image_embeddings(image_inputs, vision_model):
with torch.no_grad():
if hasattr(image_inputs, 'keys'): if hasattr(image_inputs, 'keys'):
image_inputs = {k: v.squeeze(1) if v.ndim > 2 and v.shape[1] == 1 else v for k, v in image_inputs.items()} image_inputs = {k: v.squeeze(1) if v.ndim > 2 and v.shape[1] == 1 else v for k, v in image_inputs.items()}
with torch.no_grad():
outputs = vision_model(**image_inputs) outputs = vision_model(**image_inputs)
else:
outputs = vision_model(pixel_values=image_inputs)
return outputs.last_hidden_state return outputs.last_hidden_state
@torch.compiler.disable @torch.compiler.disable