bfloat16

microsoft · Jul 22, 2024 · 88a62b1 · 88a62b1
1 parent 492096e
commit 88a62b1
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/onnxscript/ir/serde_test.py b/onnxscript/ir/serde_test.py
@@ -103,7 +103,7 @@ def test_tensor_proto_tensor_bfloat16(self):
             raw_data=raw_data,
         )
         array_from_raw_data = onnx.numpy_helper.to_array(tensor_proto_from_raw_data)
-        np.testing.assert_array_equal(array_from_raw_data, expected_array)
+        np.testing.assert_array_equal(array_from_raw_data.view(ml_dtypes.bfloat16), expected_array)
         # Test dlpack
         with self.assertRaises(BufferError):
             # NumPy does not support bfloat16 in from_dlpack