pull/130/merge
ian 2026-09-30 18:20:51 +08:00 committed by GitHub
commit 50decbdfeb
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 5 additions and 3 deletions

View File

@ -66,10 +66,12 @@ class MiniMindVLM(MiniMindForCausalLM):
@staticmethod
def get_image_embeddings(image_inputs, vision_model):
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()}
with torch.no_grad():
outputs = vision_model(**image_inputs)
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()}
outputs = vision_model(**image_inputs)
else:
outputs = vision_model(pixel_values=image_inputs)
return outputs.last_hidden_state
@torch.compiler.disable