ocr_predictor always use cpu not gpu

[iwater]

Bug description

the workload always on cpu, not gpu

Code snippet to reproduce the bug

from doctr.models import ocr_predictor
import os

os.environ["USE_TORCH"] = '1'

predictor = ocr_predictor(det_arch='db_resnet50', reco_arch='crnn_vgg16_bn', pretrained=True)

Error traceback

no error

Environment

DocTR version: 0.6.1a0 TensorFlow version: N/A PyTorch version: 2.0.1 (torchvision 0.15.2) OpenCV version: 4.7.0 OS: Ubuntu 22.04.2 LTS Python version: 3.8.16 Is CUDA available (TensorFlow): N/A Is CUDA available (PyTorch): Yes CUDA runtime version: Could not collect GPU models and configuration: GPU 0: Tesla P40 Nvidia driver version: 530.41.03 cuDNN version: Could not collect

Deep Learning backend

is_tf_available: False is_torch_available: True

[iwater]

find out works code

device = torch.device("cuda:0")
predictor = ocr_predictor(det_arch='db_resnet50', reco_arch='crnn_vgg16_bn', pretrained=True).to(device)

[sanjay-nit]

find out works code

device = torch.device("cuda:0")
predictor = ocr_predictor(det_arch='db_resnet50', reco_arch='crnn_vgg16_bn', pretrained=True).to(device)

Doesn't work in google colab AttributeError: 'OCRPredictor' object has no attribute 'to'

[oussamaHadder]

You need to set os.environ['USE_TORCH'] = 'YES' and os.environ['USE_TF'] = 'NO'. That solved the issue for me.