[Serving]ppcls preprocessor support gpu (#615)

* serving ppcls support gpu

* serving ppcls preprocessor use cpu
This commit is contained in:
heliqi
2022-11-17 17:16:32 +08:00
committed by GitHub
parent 1694d9a6a2
commit 6ebe612377
6 changed files with 17 additions and 6 deletions

View File

@@ -151,8 +151,7 @@ pybind11::capsule FDTensorToDLPack(FDTensor& fd_tensor) {
dlpack_tensor->dl_tensor.dtype = FDToDlpackType(fd_tensor.dtype);
// TODO(liqi): FDTensor add device_id
dlpack_tensor->dl_tensor.device.device_id = 0;
dlpack_tensor->dl_tensor.device.device_id = fd_tensor.device_id;
if(fd_tensor.device == Device::GPU) {
if (fd_tensor.is_pinned_memory) {
dlpack_tensor->dl_tensor.device.device_type = DLDeviceType::kDLCUDAHost;