unify to pt

mindee · Oct 12, 2023 · 3dd8580 · 3dd8580
1 parent bc46f3a
commit 3dd8580
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 2 deletions.
diff --git a/doctr/models/recognition/master/tensorflow.py b/doctr/models/recognition/master/tensorflow.py
@@ -181,7 +181,9 @@ def call(
             output = self.decoder(gt, encoded, source_mask, target_mask, **kwargs)
             logits = self.linear(output, **kwargs)
         else:
-            logits = _bf16_to_numpy_dtype(self.decode(encoded, **kwargs))
+            logits = self.decode(encoded, **kwargs)
+
+        logits = _bf16_to_numpy_dtype(logits)
 
         if self.exportable:
             out["logits"] = logits

diff --git a/doctr/models/recognition/parseq/tensorflow.py b/doctr/models/recognition/parseq/tensorflow.py
@@ -388,7 +388,9 @@ def call(
                     )
                 )
         else:
-            logits = _bf16_to_numpy_dtype(self.decode_autoregressive(features, **kwargs))
+            logits = self.decode_autoregressive(features, **kwargs)
+
+        logits = _bf16_to_numpy_dtype(logits)
 
         out: Dict[str, tf.Tensor] = {}
         if self.exportable: