commit 487b53ebb6e9d8289855ae7e5e4df09817eba6a7
author Frederic Bastien <fbastien@nvidia.com> Wed Jan 20 06:10:06 2021 -0800
committer Frederic Bastien <fbastien@nvidia.com> Wed Jan 20 06:10:06 2021 -0800
tree 861438bcbfb6530632dd3f922063837a2a165673
parent 6cc470aab0b682ebdd43580075bea16266620498

Fixup since TF internal changed since TF 2.4

tensorflow/core/common_runtime/gpu/gpu_cudamallocasync_allocator.cc

diff --git a/tensorflow/core/common_runtime/gpu/gpu_cudamallocasync_allocator.cc b/tensorflow/core/common_runtime/gpu/gpu_cudamallocasync_allocator.cc
index fa6fe57..d9c10ab 100644
--- a/tensorflow/core/common_runtime/gpu/gpu_cudamallocasync_allocator.cc
+++ b/tensorflow/core/common_runtime/gpu/gpu_cudamallocasync_allocator.cc
@@ -21,8 +21,8 @@
 
 #include "tensorflow/core/common_runtime/gpu/gpu_cudamallocasync_allocator.h"
 
+#include "tensorflow/core/common_runtime/device/device_id_utils.h"
 #include "tensorflow/core/common_runtime/gpu/gpu_id.h"
-#include "tensorflow/core/common_runtime/gpu/gpu_id_utils.h"
 #include "tensorflow/core/common_runtime/gpu/gpu_init.h"
 #include "tensorflow/core/framework/allocator.h"
 #include "tensorflow/core/platform/stream_executor.h"
@@ -37,7 +37,8 @@
       name_(absl::StrCat("gpu_async_", platform_gpu_id.value())),
       pool_(nullptr) {
   stream_exec_ =
-      GpuIdUtil::ExecutorForPlatformGpuId(platform_gpu_id).ValueOrDie();
+      DeviceIdUtil::ExecutorForPlatformDeviceId(GPUMachineManager(),
+                                                platform_gpu_id).ValueOrDie();
 
 #if CUDA_VERSION < 11020
   LOG(FATAL) << "TF_GPU_ALLOCATOR=cuda_malloc_async need CUDA 11.2 or higher to compile.";