set device_id in eager_gen

PaddlePaddle · Apr 6, 2022 · 3537f12 · 3537f12 · paddle-bot-old · Apr 6, 2022
1 parent a0835a1
commit 3537f12
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 10 deletions.
diff --git a/paddle/fluid/eager/auto_code_generator/final_state_generator/eager_gen.py b/paddle/fluid/eager/auto_code_generator/final_state_generator/eager_gen.py
@@ -195,6 +195,16 @@ class {} : public egr::GradNodeBase {{
     
     // Get Input AutoGradMeta
 {}
+    // Set Device Id
+    auto place = egr::Controller::Instance().GetExpectedPlace();
+    if (paddle::platform::is_gpu_place(place)) {{
+#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
+      phi::backends::gpu::SetDeviceId(place.device);
+#else
+      PADDLE_THROW(paddle::platform::errors::PreconditionNotMet(
+        "PaddlePaddle should compile with GPU if use CUDAPlace."));
+#endif
+    }}
     // Forward API Call
 {}
     // Get Outputs
@@ -285,6 +295,7 @@ class {} : public egr::GradNodeBase {{
 #include "paddle/fluid/platform/profiler/event_tracing.h"
 #include "paddle/fluid/eager/amp_utils.h"
 #include "paddle/fluid/eager/eager_amp_auto_cast.h"
+#include "paddle/phi/backends/gpu/gpu_info.h"
 
 {}
 {}

diff --git a/paddle/fluid/imperative/tracer.cc b/paddle/fluid/imperative/tracer.cc
@@ -385,16 +385,6 @@ void Tracer::TraceOp(const std::string& type, const NameTensorMap& ins,
 
 void Tracer::SetExpectedPlace(platform::Place place) {
   expected_place_ = place;
-  if (platform::is_gpu_place(place)) {
-#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
-    if (place.device) {
-      platform::SetDeviceId(place.device);
-    }
-#else
-    PADDLE_THROW(platform::errors::PreconditionNotMet(
-        "PaddlePaddle should compile with GPU if use CUDAPlace."));
-#endif
-  }
 }
 
 bool Tracer::ComputeRequiredGrad(const NameVarBaseMap& ins,