From e942c16e65fe6352d5d1608aee21b66b4ec36c50 Mon Sep 17 00:00:00 2001 From: felix Date: Thu, 12 Oct 2023 08:43:29 +0200 Subject: [PATCH] typing and order --- doctr/models/utils/pytorch.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/doctr/models/utils/pytorch.py b/doctr/models/utils/pytorch.py index 4e15fa628a..2179ba572b 100644 --- a/doctr/models/utils/pytorch.py +++ b/doctr/models/utils/pytorch.py @@ -25,6 +25,11 @@ def _copy_tensor(x: torch.Tensor) -> torch.Tensor: return x.clone().detach() +def _bf16_to_numpy_dtype(x: torch.Tensor) -> torch.Tensor: + # bfloat16 is not supported in .numpy(): torch/csrc/utils/tensor_numpy.cpp:aten_to_numpy_dtype + return x.float() if x.dtype == torch.bfloat16 else x + + def load_pretrained_params( model: nn.Module, url: Optional[str] = None, @@ -157,8 +162,3 @@ def export_model_to_onnx(model: nn.Module, model_name: str, dummy_input: torch.T ) logging.info(f"Model exported to {model_name}.onnx") return f"{model_name}.onnx" - - -def _bf16_to_numpy_dtype(x): - # bfloat16 is not supported in .numpy(): torch/csrc/utils/tensor_numpy.cpp:aten_to_numpy_dtype - return x.float() if x.dtype == torch.bfloat16 else x