x as dict

2021-07-06 07:55:05 +00:00 · 2021-07-06 07:55:05 +00:00 · 80561b154e
parent 202a0b5b85
commit 80561b154e
1 changed files with 4 additions and 1 deletions
--- 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: