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))
|