Use the NVML from the CUDA.

Google-ML-Automation · Google-ML-Automation · commit 6f7d02a86c3b · 2025-03-13T00:21:16.000-07:00
PiperOrigin-RevId: 736402658
diff --git a/xla/service/gpu/model/BUILD b/xla/service/gpu/model/BUILD
@@ -397,7 +397,7 @@ cc_library(
         "@llvm-project//llvm:Support",
         "@llvm-project//mlir:IR",
         "@tsl//tsl/platform:status",
-    ] + if_cuda_is_configured(["@local_config_cuda//cuda:nvml_headers"]),
+    ] + if_cuda_is_configured(["@local_config_cuda//cuda:cuda_headers"]),
 )
 
 xla_cc_test(

Original file line number	Diff line number	Diff line change
`@@ -397,7 +397,7 @@ cc_library(`
`397`	`397`	`"@llvm-project//llvm:Support",`
`398`	`398`	`"@llvm-project//mlir:IR",`
`399`	`399`	`"@tsl//tsl/platform:status",`
`400`		`- ] + if_cuda_is_configured(["@local_config_cuda//cuda:nvml_headers"]),`
	`400`	`+ ] + if_cuda_is_configured(["@local_config_cuda//cuda:cuda_headers"]),`
`401`	`401`	`)`
`402`	`402`
`403`	`403`	`xla_cc_test(`