change type check

doctr/models/classification/zoo.py

@@ -50,10 +50,9 @@ def _orientation_predictor(
     else:
         allowed_archs = [classification.MobileNetV3]
         if is_torch_available():
-            # The following is required for torch compiled models
-            import torch
+            from doctr.models.utils import _get_torch_compile_type
 
-            allowed_archs.append(torch._dynamo.eval_frame.OptimizedModule)
+            allowed_archs.append(_get_torch_compile_type())
 
         if not isinstance(arch, tuple(allowed_archs)):
             raise ValueError(f"unknown architecture: {type(arch)}")

doctr/models/detection/zoo.py

@@ -59,9 +59,9 @@ def _predictor(arch: Any, pretrained: bool, assume_straight_pages: bool = True,
         allowed_archs = [detection.DBNet, detection.LinkNet, detection.FAST]
         if is_torch_available():
             # The following is required for torch compiled models
-            import torch
+            from doctr.models.utils import _get_torch_compile_type
 
-            allowed_archs.append(torch._dynamo.eval_frame.OptimizedModule)
+            allowed_archs.append(_get_torch_compile_type())
 
         if not isinstance(arch, tuple(allowed_archs)):
             raise ValueError(f"unknown architecture: {type(arch)}")

doctr/models/recognition/zoo.py

@@ -38,9 +38,9 @@ def _predictor(arch: Any, pretrained: bool, **kwargs: Any) -> RecognitionPredict
         allowed_archs = [recognition.CRNN, recognition.SAR, recognition.MASTER, recognition.ViTSTR, recognition.PARSeq]
         if is_torch_available():
             # The following is required for torch compiled models
-            import torch
+            from doctr.models.utils import _get_torch_compile_type
 
-            allowed_archs.append(torch._dynamo.eval_frame.OptimizedModule)
+            allowed_archs.append(_get_torch_compile_type())
 
         if not isinstance(arch, tuple(allowed_archs)):
             raise ValueError(f"unknown architecture: {type(arch)}")

doctr/models/utils/pytorch.py

@@ -18,9 +18,14 @@
     "export_model_to_onnx",
     "_copy_tensor",
     "_bf16_to_float32",
+    "_get_torch_compile_type",
 ]
 
 
+def _get_torch_compile_type() -> Any:
+    return torch._dynamo.eval_frame.OptimizedModule
+
+
 def _copy_tensor(x: torch.Tensor) -> torch.Tensor:
     return x.clone().detach()