Make cuda_dnn use StreamExecutor::DeviceMemoryUsage rather than GpuDr…

…iver API to get the same thing. PiperOrigin-RevId: 684603618
openxla · Oct 10, 2024 · 704b1d7 · 704b1d7
1 parent 1ffcdaa
commit 704b1d7
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/xla/stream_executor/cuda/cuda_dnn.cc b/xla/stream_executor/cuda/cuda_dnn.cc
@@ -534,7 +534,7 @@ absl::Status CudnnSupport::Init() {
   LOG(ERROR) << "Could not create cudnn handle: "
              << CudnnStatusToString(status);
   int64_t free, total;
-  GpuDriver::GetDeviceMemoryInfo(parent_->gpu_context(), &free, &total);
+  parent_->DeviceMemoryUsage(&free, &total);
   LOG(ERROR) << "Memory usage: " << free << " bytes free, " << total
              << " bytes total.";