Skip to content

Commit f6c80c6

Browse files
[bug] Fix missing model check
1 parent e87736e commit f6c80c6

File tree

1 file changed

+9
-1
lines changed
  • onnxtr/models/recognition

1 file changed

+9
-1
lines changed

onnxtr/models/recognition/zoo.py

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -36,7 +36,15 @@ def _predictor(
3636
_model = recognition.__dict__[arch](load_in_8_bit=load_in_8_bit, engine_cfg=engine_cfg)
3737
else:
3838
if not isinstance(
39-
arch, (recognition.CRNN, recognition.SAR, recognition.MASTER, recognition.ViTSTR, recognition.PARSeq)
39+
arch,
40+
(
41+
recognition.CRNN,
42+
recognition.SAR,
43+
recognition.MASTER,
44+
recognition.ViTSTR,
45+
recognition.PARSeq,
46+
recognition.VIPTR,
47+
),
4048
):
4149
raise ValueError(f"unknown architecture: {type(arch)}")
4250
_model = arch

0 commit comments

Comments
 (0)