unknown architecture: <class 'torch._dynamo.eval_frame.OptimizedModule'>

[VMM-MMV]

Bug description

I am running this exact code from https://mindee.github.io/doctr/latest/using_doctr/using_model_export.html

But I am getting this error.

Code snippet to reproduce the bug

from doctr.models import (
    ocr_predictor,
    vitstr_small,
    fast_base,
    mobilenet_v3_small_crop_orientation,
    mobilenet_v3_small_page_orientation,
    crop_orientation_predictor,
    page_orientation_predictor
)

# Compile the models
detection_model = torch.compile(
    fast_base(pretrained=True).eval()
)
recognition_model = torch.compile(
    vitstr_small(pretrained=True).eval()
)
crop_orientation_model = torch.compile(
    mobilenet_v3_small_crop_orientation(pretrained=True).eval()
)
page_orientation_model = torch.compile(
    mobilenet_v3_small_page_orientation(pretrained=True).eval()
)

predictor = models.ocr_predictor(
    detection_model, recognition_model, assume_straight_pages=False
)
# NOTE: Only required for non-straight pages (`assume_straight_pages=False`) and non-disabled orientation classification
# Set the orientation predictors
predictor.crop_orientation_predictor = crop_orientation_predictor(crop_orientation_model)
predictor.page_orientation_predictor = page_orientation_predictor(page_orientation_model)

Error traceback

Traceback (most recent call last):
  File "/home/serveruser/groq/code/utils/ocr.py", line 32, in <module>
    predictor = ocr_predictor(
  File "/home/serveruser/anaconda3/envs/classify/lib/python3.10/site-packages/doctr/models/zoo.py", line 114, in ocr_predictor
    return _predictor(
  File "/home/serveruser/anaconda3/envs/classify/lib/python3.10/site-packages/doctr/models/zoo.py", line 32, in _predictor
    det_predictor = detection_predictor(
  File "/home/serveruser/anaconda3/envs/classify/lib/python3.10/site-packages/doctr/models/detection/zoo.py", line 103, in detection_predictor
    return _predictor(arch, pretrained, assume_straight_pages, **kwargs)
  File "/home/serveruser/anaconda3/envs/classify/lib/python3.10/site-packages/doctr/models/detection/zoo.py", line 60, in _predictor
    raise ValueError(f"unknown architecture: {type(arch)}")
ValueError: unknown architecture: <class 'torch._dynamo.eval_frame.OptimizedModule'>

Environment

DocTR version: v0.10.0
TensorFlow version: N/A
PyTorch version: 2.6.0.dev20241216+cu118 (torchvision 0.22.0.dev20241216+cu118)
OpenCV version: 4.10.0
OS: Ubuntu 24.04.1 LTS
Python version: 3.10.15
Is CUDA available (TensorFlow): N/A
Is CUDA available (PyTorch): Yes
CUDA runtime version: Could not collect
GPU models and configuration: GPU 0: NVIDIA GeForce RTX 3060
Nvidia driver version: 565.57.01
cuDNN version: Could not collect

(classify) serveruser@server:~/groq$ nvidia-smi
+-----------------------------------------------------------------------------------------+
| NVIDIA-SMI 565.57.01 Driver Version: 565.57.01 CUDA Version: 12.7 |
|-----------------------------------------+------------------------+----------------------+

Deep Learning backend

is_tf_available: False
is_torch_available: True

[felixdittrich92]

Hi @VMM-MMV 👋

You use our latest doc status which matches the dev version 0.11.0a0 but you have installed the latest pypi version v0.10 :)

If you want to test the torch.compile feature please install the repo from the main branch:

For example:

pip3 install python-doctr[torch,viz]@git+https://github.com/mindee/doctr.git

The feature will be official available with our next release (v0.11).

Best,
Felix

[VMM-MMV]

Thanks.

Appreciate the work you do!