From fb0cadfd2fa159a3d949357300a668a9cff75802 Mon Sep 17 00:00:00 2001 From: From00 Date: Wed, 2 Mar 2022 10:05:45 +0800 Subject: [PATCH] Fix bug for prepare phi OP (#40033) --- paddle/fluid/imperative/prepared_operator.cc | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/paddle/fluid/imperative/prepared_operator.cc b/paddle/fluid/imperative/prepared_operator.cc index 9dd1dacc02c..2317bfdd7c0 100644 --- a/paddle/fluid/imperative/prepared_operator.cc +++ b/paddle/fluid/imperative/prepared_operator.cc @@ -186,11 +186,10 @@ PreparedOp PrepareImpl(const NameVarMap& ins, << " | kernel key: " << pt_kernel_key << " | kernel: " << pt_kernel; - if (platform::is_cpu_place(expected_kernel_key.place_)) { - auto* cpu_ctx = pool.Get(paddle::platform::CPUPlace()); - return PreparedOp(op, ctx, expected_kernel_key, pt_kernel_signature, - pt_kernel, cpu_ctx); + if (expected_kernel_key.place_ != place) { + dev_ctx = pool.Get(expected_kernel_key.place_); } + // TODO(chenweihang): using CPUKernel when miss device kernel case return PreparedOp(op, ctx, expected_kernel_key, pt_kernel_signature, pt_kernel, dev_ctx); -- GitLab