Update deviceimpl.cpp

DeepLink-org · Mar 31, 2024 · 0f98da3 · 0f98da3
1 parent 99bd01c
commit 0f98da3
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/dipu/torch_dipu/csrc_dipu/vendor/ascend/deviceimpl.cpp b/dipu/torch_dipu/csrc_dipu/vendor/ascend/deviceimpl.cpp
@@ -28,7 +28,10 @@ void finalizeVendor() { DIPU_CALLACLRT(aclFinalize()); }
 
 deviceId_t current_device() {
   if (currentDeviceIndex < 0) {
-    setDevice(-1);  // need fix......
+    int device_index = 0;
+    DIPU_CALLACLRT(::aclrtGetDevice(&device_index))
+    return device_index;
+    // setDevice(-1);  // need fix......
     // DIPU_CALLACLRT(::aclrtGetDevice(&currentDeviceIndex))
   }
   return static_cast<deviceId_t>(currentDeviceIndex);