test: Added PyTorch zoo unittest

mindee · Sep 29, 2021 · 110b8c5 · 110b8c5
1 parent f14296d
commit 110b8c5
Showing 1 changed file with 17 additions and 0 deletions.
diff --git a/test/pytorch/test_models_zoo_pt.py b/test/pytorch/test_models_zoo_pt.py
@@ -0,0 +1,17 @@
+import pytest
+
+from doctr import models
+from doctr.models.predictor import OCRPredictor
+
+
+@pytest.mark.parametrize(
+    "det_arch, reco_arch",
+    [
+        ["db_mobilenet_v3_large", "crnn_mobilenet_v3_large"],
+    ],
+)
+def test_zoo_models(det_arch, reco_arch):
+    # Model
+    predictor = models.ocr_predictor(det_arch, reco_arch, pretrained=True)
+    # Output checks
+    assert isinstance(predictor, OCRPredictor)