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) {