fix: fix the bug when return_numpy is false
Browse files
modeling_jina_embeddings_v4.py
CHANGED
@@ -356,9 +356,11 @@ class JinaEmbeddingsV4Model(Qwen2_5_VLForConditionalGeneration):
|
|
356 |
# Remove padding by selecting only valid tokens for each sequence
|
357 |
embeddings = [emb[mask] for emb, mask in zip(embeddings, valid_tokens)]
|
358 |
# Stack back into tensor with variable sequence lengths
|
359 |
-
results.
|
360 |
else:
|
361 |
results.append(
|
|
|
|
|
362 |
embeddings.cpu()
|
363 |
if return_numpy
|
364 |
else list(torch.unbind(embeddings))
|
|
|
356 |
# Remove padding by selecting only valid tokens for each sequence
|
357 |
embeddings = [emb[mask] for emb, mask in zip(embeddings, valid_tokens)]
|
358 |
# Stack back into tensor with variable sequence lengths
|
359 |
+
results.append(embeddings)
|
360 |
else:
|
361 |
results.append(
|
362 |
+
# If return_numpy is True, move embeddings to CPU for numpy conversion
|
363 |
+
# Otherwise, unbind the tensor into a list of individual tensors along dim=0
|
364 |
embeddings.cpu()
|
365 |
if return_numpy
|
366 |
else list(torch.unbind(embeddings))
|