[Bug] Fixed assume_straight_pages for custom models (#1681)

mindee · Jul 29, 2024 · a659425 · a659425
1 parent 1cea7d8
commit a659425
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/doctr/models/detection/zoo.py b/doctr/models/detection/zoo.py
@@ -61,6 +61,7 @@ def _predictor(arch: Any, pretrained: bool, assume_straight_pages: bool = True,
 
         _model = arch
         _model.assume_straight_pages = assume_straight_pages
+        _model.postprocessor.assume_straight_pages = assume_straight_pages
 
     kwargs.pop("pretrained_backbone", None)