[BugFix]Ensure that bf16 arrays are created as expected (#16436)

Co-authored-by: Bin Li <[email protected]>
apache · Jan 20, 2024 · ccca00a · ccca00a
1 parent ffa404f
commit ccca00a
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/python/tvm/runtime/ndarray.py b/python/tvm/runtime/ndarray.py
@@ -176,6 +176,8 @@ def copyfrom(self, source_array):
         if (not source_array.flags["C_CONTIGUOUS"]) or (
             dtype == "bfloat16" or dtype != np_dtype_str
         ):
+            if dtype == "bfloat16":
+                source_array = np.frombuffer(source_array.tobytes(), "uint16")
             source_array = np.ascontiguousarray(
                 source_array, dtype="uint16" if dtype == "bfloat16" else dtype
             )