jina-ai / executors

internal-only
Apache License 2.0
31 stars 12 forks source link

feat(gpu): split ranker into cpu and gpu #198

Closed jacobowitz closed 2 years ago

jacobowitz commented 2 years ago

This is part of https://github.com/jina-ai/executors/issues/169

It splits the DPRReaderRanker into cpu/gpu versions. Also I noticed that the GPU test was actually failing, because the Tokenizer was generating tensors always on CPU device. With this PR it can also generate GPU based tensors. The results always have to be copied to CPU then to further process them on the CPU