From 7f5009c1be820ba7a98f17765e8aff289ab445e7 Mon Sep 17 00:00:00 2001
From: bxia <2608468143@qq.com>
Date: Mon, 25 Sep 2023 09:00:36 +0800
Subject: [PATCH] Adapt to rocm

---
 setup.py | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/setup.py b/setup.py
index 45e2e096d6..0b0e2ecfbe 100644
--- a/setup.py
+++ b/setup.py
@@ -259,6 +259,9 @@ def get_extensions():
             define_macros += [('MMCV_WITH_CUDA', None)]
             cuda_args = os.getenv('MMCV_CUDA_ARGS')
             extra_compile_args['nvcc'] = [cuda_args] if cuda_args else []
+            if is_rocm_pytorch and platform.system() != 'Windows':
+                extra_compile_args['nvcc'] += \
+                    ['--gpu-max-threads-per-block=1024']
             op_files = glob.glob('./mmcv/ops/csrc/pytorch/*.cpp') + \
                 glob.glob('./mmcv/ops/csrc/pytorch/cpu/*.cpp') + \
                 glob.glob('./mmcv/ops/csrc/pytorch/cuda/*.cu') + \