x as dict

80561b15 · LDOUBLEV · 202a0b5b · 80561b15
Commit 80561b15 authored Jul 06, 2021 by LDOUBLEV
Hide whitespace changes
Inline Side-by-side

Showing with 4 additions and 1 deletion

ppocr/modeling/architectures/base_model.py ppocr/modeling/architectures/base_model.py +4 -1

No files found.
--- a/ppocr/modeling/architectures/base_model.py
+++ b/ppocr/modeling/architectures/base_model.py
@@ -79,7 +79,10 @@ class BaseModel(nn.Layer):
            x = self.neck(x)
        y["neck_out"] = x
        x = self.head(x, targets=data)
-        y["head_out"] = x
+        if type(x) is dict:
+            y.update(x)
+        else:
+            y["head_out"] = x
        if self.return_all_feats:
            return y
        else: