update type check + mypy torch decorator

doctr/models/classification/zoo.py

@@ -48,13 +48,14 @@ def _orientation_predictor(
         # Load directly classifier from backbone
         _model = classification.__dict__[arch](pretrained=pretrained)
     else:
-        allowed_archs = (classification.MobileNetV3,)
+        allowed_archs = [classification.MobileNetV3]
         if is_torch_available():
+            # The following is required for torch compiled models
             import torch
 
-            allowed_archs += (torch._dynamo.eval_frame.OptimizedModule,)
+            allowed_archs.append(torch._dynamo.eval_frame.OptimizedModule)
 
-        if not isinstance(arch, allowed_archs):
+        if not isinstance(arch, tuple(allowed_archs)):
             raise ValueError(f"unknown architecture: {type(arch)}")
         _model = arch
 

doctr/models/detection/differentiable_binarization/pytorch.py

@@ -205,7 +205,7 @@ def forward(
 
         if target is None or return_preds:
             # Disable for torch.compile compatibility
-            @torch.compiler.disable
+            @torch.compiler.disable  # type: ignore[attr-defined]
             def _postprocess(prob_map: torch.Tensor) -> List[Dict[str, Any]]:
                 return [
                     dict(zip(self.class_names, preds))

doctr/models/detection/fast/pytorch.py

@@ -196,7 +196,7 @@ def forward(
 
         if target is None or return_preds:
             # Disable for torch.compile compatibility
-            @torch.compiler.disable
+            @torch.compiler.disable  # type: ignore[attr-defined]
             def _postprocess(prob_map: torch.Tensor) -> List[Dict[str, Any]]:
                 return [
                     dict(zip(self.class_names, preds))

doctr/models/detection/linknet/pytorch.py

@@ -183,7 +183,7 @@ def forward(
 
         if target is None or return_preds:
             # Disable for torch.compile compatibility
-            @torch.compiler.disable
+            @torch.compiler.disable  # type: ignore[attr-defined]
             def _postprocess(prob_map: torch.Tensor) -> List[Dict[str, Any]]:
                 return [
                     dict(zip(self.class_names, preds))

doctr/models/detection/zoo.py

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

doctr/models/recognition/crnn/pytorch.py

@@ -213,7 +213,7 @@ def forward(
 
         if target is None or return_preds:
             # Disable for torch.compile compatibility
-            @torch.compiler.disable
+            @torch.compiler.disable  # type: ignore[attr-defined]
             def _postprocess(logits: torch.Tensor) -> List[Tuple[str, float]]:
                 return self.postprocessor(logits)
 

doctr/models/recognition/master/pytorch.py

@@ -209,7 +209,7 @@ def forward(
 
         if return_preds:
             # Disable for torch.compile compatibility
-            @torch.compiler.disable
+            @torch.compiler.disable  # type: ignore[attr-defined]
             def _postprocess(logits: torch.Tensor) -> List[Tuple[str, float]]:
                 return self.postprocessor(logits)
 

doctr/models/recognition/parseq/pytorch.py

@@ -372,7 +372,7 @@ def forward(
 
         if target is None or return_preds:
             # Disable for torch.compile compatibility
-            @torch.compiler.disable
+            @torch.compiler.disable  # type: ignore[attr-defined]
             def _postprocess(logits: torch.Tensor) -> List[Tuple[str, float]]:
                 return self.postprocessor(logits)
 

doctr/models/recognition/sar/pytorch.py

@@ -262,7 +262,7 @@ def forward(
 
         if target is None or return_preds:
             # Disable for torch.compile compatibility
-            @torch.compiler.disable
+            @torch.compiler.disable  # type: ignore[attr-defined]
             def _postprocess(decoded_features: torch.Tensor) -> List[Tuple[str, float]]:
                 return self.postprocessor(decoded_features)
 

doctr/models/recognition/vitstr/pytorch.py

@@ -107,7 +107,7 @@ def forward(
 
         if target is None or return_preds:
             # Disable for torch.compile compatibility
-            @torch.compiler.disable
+            @torch.compiler.disable  # type: ignore[attr-defined]
             def _postprocess(decoded_features: torch.Tensor) -> List[Tuple[str, float]]:
                 return self.postprocessor(decoded_features)
 

doctr/models/recognition/zoo.py

@@ -35,12 +35,14 @@ def _predictor(arch: Any, pretrained: bool, **kwargs: Any) -> RecognitionPredict
             pretrained=pretrained, pretrained_backbone=kwargs.get("pretrained_backbone", True)
         )
     else:
-        allowed_archs = (recognition.CRNN, recognition.SAR, recognition.MASTER, recognition.ViTSTR, recognition.PARSeq)
+        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
 
-            allowed_archs += (torch._dynamo.eval_frame.OptimizedModule,)
-        if not isinstance(arch, allowed_archs):
+            allowed_archs.append(torch._dynamo.eval_frame.OptimizedModule)
+
+        if not isinstance(arch, tuple(allowed_archs)):
             raise ValueError(f"unknown architecture: {type(arch)}")
         _model = arch