diff --git a/python/pyarrow/tensorflow/plasma_op.cc b/python/pyarrow/tensorflow/plasma_op.cc index 2083e57c42f3d..3b7fe2a7a13ad 100644 --- a/python/pyarrow/tensorflow/plasma_op.cc +++ b/python/pyarrow/tensorflow/plasma_op.cc @@ -189,7 +189,7 @@ class TensorToPlasmaOp : public tf::AsyncOpKernel { // async memcpy. Under the hood it performs cuMemHostRegister(), see: // http://docs.nvidia.com/cuda/cuda-driver-api/group__CUDA__MEM.html#group__CUDA__MEM_1gf0a9fe11544326dabd743b7aa6b54223 CHECK(stream_executor->HostMemoryRegister(static_cast(data), - static_cast(total_bytes))); + static_cast(total_bytes))); { tf::mutex_lock l(d2h_stream_mu);