enable rocm, fix transform.h

2023-03-08 06:37:34 +01:00 · 2023-03-08 06:37:34 +01:00 · 62c4efac51
commit 62c4efac51
parent ba9e00d911
1 changed files with 6 additions and 0 deletions
--- 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) {