Use array interface in Python prediction return. (#9855)

This commit is contained in:
Jiaming Yuan
2023-12-08 03:42:14 +08:00
committed by GitHub
parent 2c0fc97306
commit 39c637ee19
4 changed files with 49 additions and 48 deletions

View File

@@ -683,7 +683,7 @@ using MatrixView = TensorView<T, 2>;
*
* `stream` is optionally included when data is on CUDA device.
*/
template <typename T, int32_t D>
template <typename T, std::int32_t D>
Json ArrayInterface(TensorView<T const, D> const &t) {
Json array_interface{Object{}};
array_interface["data"] = std::vector<Json>(2);
@@ -691,7 +691,7 @@ Json ArrayInterface(TensorView<T const, D> const &t) {
array_interface["data"][1] = Boolean{true};
if (t.Device().IsCUDA()) {
// Change this once we have different CUDA stream.
array_interface["stream"] = Null{};
array_interface["stream"] = Integer{2};
}
std::vector<Json> shape(t.Shape().size());
std::vector<Json> stride(t.Stride().size());