diff --git a/paddle/fluid/operators/conditional_block_op.cc b/paddle/fluid/operators/conditional_block_op.cc index 337b34e8f0bf4cb89753235205be9eb058dd01ab..bbe297206ebac87c4f88f0b2afa02ad809d59d32 100644 --- a/paddle/fluid/operators/conditional_block_op.cc +++ b/paddle/fluid/operators/conditional_block_op.cc @@ -54,7 +54,18 @@ class ConditionalOp : public framework::OperatorBase { "numel should be 1, actual numel is %d", ips[0]->numel()); } - return ips[0]->data()[0]; + bool res; + if (platform::is_gpu_place(ips[0]->place())) { +#ifdef PADDLE_WITH_CUDA + framework::LoDTensor cpu_tensor; + framework::TensorCopy(*ips[0], platform::CPUPlace(), &cpu_tensor); + platform::DeviceContextPool::Instance().Get(ips[0]->place())->Wait(); + res = cpu_tensor.data()[0]; +#endif + } else { + res = ips[0]->data()[0]; + } + return res; } };