Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[tokenizer] Ensure GPU is used in TextEmbeddingTranslator #3212

Conversation

david-sitsky
Copy link
Contributor

This fixes #3211, where using djl://ai.djl.huggingface.pytorch/sentence-transformers/clip-ViT-B-32-multilingual-v1 on 0.28.0 fails on a GPU machine with:

Caused by: ai.djl.engine.EngineException: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu! (when checking argument for argument mat2 in method wrapper_mm)
	at ai.djl.pytorch.jni.PyTorchLibrary.torchNNLinear(PyTorchLibrary.java) ~[pytorch-engine-0.28.0.jar:?]
	at ai.djl.pytorch.jni.JniUtils.linear(JniUtils.java:1376) ~[pytorch-engine-0.28.0.jar:?]
	at ai.djl.pytorch.engine.PtNDArrayEx.linear(PtNDArrayEx.java:397) ~[pytorch-engine-0.28.0.jar:?]
	at ai.djl.huggingface.translator.TextEmbeddingTranslator.processEmbedding(TextEmbeddingTranslator.java:181) ~[tokenizers-0.28.0.jar:?]
	at ai.djl.huggingface.translator.TextEmbeddingTranslator.batchProcessOutput(TextEmbeddingTranslator.java:144) ~[tokenizers-0.28.0.jar:?]
	at ai.djl.inference.Predictor.batchPredict(Predictor.java:190) ~[api-0.28.0.jar:?]
	at ai.djl.inference.Predictor.predict(Predictor.java:132) ~[api-0.28.0.jar:?]

@david-sitsky david-sitsky requested review from zachgk, frankfliu and a team as code owners May 22, 2024 07:20
@frankfliu frankfliu merged commit d3fd000 into deepjavalibrary:master May 22, 2024
5 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

TextEmbeddingTranslator fails with "EngineException: Expected all tensors to be on the same device"
2 participants