Skip to content

Commit

Permalink
[tokenizer] Ensure GPU is used in TextEmbeddingTranslator (#3212)
Browse files Browse the repository at this point in the history
* Ensure GPU is used if the manager is configured that way.

---------

Co-authored-by: David Sitsky <sits@nuix.com>
Co-authored-by: Frank Liu <frankfliu2000@gmail.com>
  • Loading branch information
3 people authored May 22, 2024
1 parent a21baf6 commit d3fd000
Showing 1 changed file with 1 addition and 2 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,6 @@
*/
package ai.djl.huggingface.translator;

import ai.djl.Device;
import ai.djl.huggingface.tokenizers.Encoding;
import ai.djl.huggingface.tokenizers.HuggingFaceTokenizer;
import ai.djl.ndarray.NDArray;
Expand Down Expand Up @@ -78,7 +77,7 @@ public Batchifier getBatchifier() {
/** {@inheritDoc} */
@Override
public void prepare(TranslatorContext ctx) throws Exception {
NDManager manager = ctx.getPredictorManager().newSubManager(Device.cpu());
NDManager manager = ctx.getPredictorManager().newSubManager();
if (dense != null) {
Path file = Paths.get(dense);
if (!file.isAbsolute()) {
Expand Down

0 comments on commit d3fd000

Please sign in to comment.