diff --git a/src/c_api/c_api.cu b/src/c_api/c_api.cu
index 61e6ca44e..89830b89b 100644
--- a/src/c_api/c_api.cu
+++ b/src/c_api/c_api.cu
@@ -50,13 +50,21 @@ void XGBBuildInfoDevice(Json *p_info) {
 void XGBoostAPIGuard::SetGPUAttribute() {
   // Not calling `safe_cuda` to avoid unnecessary exception handling overhead.
   // If errors, do nothing, assuming running on CPU only machine.
+#if defined(XGBOOST_USE_CUDA)
   cudaGetDevice(&device_id_);
+#elif defined(XGBOOST_USE_HIP)
+  hipGetDevice(&device_id_);
+#endif
 }
 
 void XGBoostAPIGuard::RestoreGPUAttribute() {
   // Not calling `safe_cuda` to avoid unnecessary exception handling overhead.
   // If errors, do nothing, assuming running on CPU only machine.
+#if defined(XGBOOST_USE_CUDA)
   cudaSetDevice(device_id_);
+#elif defined(XGBOOST_USE_HIP)
+  hipSetDevice(device_id_);
+#endif
 }
 }                        // namespace xgboost