From 284a3edfec086bd87696f27564afc3820b26c8d8 Mon Sep 17 00:00:00 2001 From: Xuchun Shang Date: Tue, 16 Dec 2025 15:26:16 +0800 Subject: [PATCH] [bugfix] fix bfloat16 for get_tensor (#1216) * [bugfix] fix bfloat16 for get_tensor Signed-off-by: Xuchun Shang * update Signed-off-by: Xuchun Shang * update Signed-off-by: Xuchun Shang --------- Signed-off-by: Xuchun Shang --- mooncake-integration/store/store_py.cpp | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/mooncake-integration/store/store_py.cpp b/mooncake-integration/store/store_py.cpp index 1c8d45d4..4815553e 100644 --- a/mooncake-integration/store/store_py.cpp +++ b/mooncake-integration/store/store_py.cpp @@ -378,6 +378,11 @@ class MooncakeStorePyWrapper { } pybind11::object tensor = torch_module().attr("from_numpy")(np_array); + if (dtype_enum == TensorDtype::BFLOAT16) { + tensor = tensor.attr("view")(torch_module().attr("bfloat16")); + } else if (dtype_enum == TensorDtype::FLOAT16) { + tensor = tensor.attr("view")(torch_module().attr("float16")); + } return tensor; } catch (const pybind11::error_already_set &e) { @@ -494,6 +499,13 @@ class MooncakeStorePyWrapper { np_array = np_array.attr("reshape")(shape_tuple); } pybind11::object tensor = torch.attr("from_numpy")(np_array); + if (dtype_enum == TensorDtype::BFLOAT16) { + tensor = + tensor.attr("view")(torch_module().attr("bfloat16")); + } else if (dtype_enum == TensorDtype::FLOAT16) { + tensor = + tensor.attr("view")(torch_module().attr("float16")); + } results_list.append(tensor); } } catch (const pybind11::error_already_set &e) {