diff --git a/src/common/transform.h b/src/common/transform.h
index 5f9c3f1bf..974ee86d6 100644
--- a/src/common/transform.h
+++ b/src/common/transform.h
@@ -140,7 +140,13 @@ class Transform {
       // granularity is used in data vector.
       size_t shard_size = range_size;
       Range shard_range {0, static_cast<Range::DifferenceType>(shard_size)};
+
+#if defined(XGBOOST_USE_HIP)
+      dh::safe_cuda(hipSetDevice(device_));
+#else
       dh::safe_cuda(cudaSetDevice(device_));
+#endif
+
       const int kGrids =
           static_cast<int>(DivRoundUp(*(range_.end()), kBlockThreads));
       if (kGrids == 0) {