fix GET_THREADS() for ROCm (#2997)

dbdc7267 · Ashish Farmer · GitHub · 8e878f0f · dbdc7267
Unverified Commit dbdc7267 authored Nov 16, 2020 by Ashish Farmer Committed by GitHub Nov 16, 2020
Hide whitespace changes
Inline Side-by-side

Showing with 3 additions and 0 deletions

torchvision/csrc/cuda/DeformConv_cuda.cu torchvision/csrc/cuda/DeformConv_cuda.cu +3 -0

No files found.
--- a/torchvision/csrc/cuda/DeformConv_cuda.cu
+++ b/torchvision/csrc/cuda/DeformConv_cuda.cu
@@ -81,6 +81,9 @@
 const int kMaxParallelImgs = 32;

 inline unsigned int GET_THREADS() {
+#ifdef __HIP_PLATFORM_HCC__
+  return 256;
+#endif
  if (at::cuda::getCurrentDeviceProperties()->major >= 6) {
    return 1024;
  }