| 1234567891011121314 |
- import paddle
- from paddle.utils.cpp_extension import CppExtension, CUDAExtension, setup
- if __name__ == "__main__":
- if paddle.device.is_compiled_with_cuda():
- setup(
- name='rbox_iou_ops',
- ext_modules=CUDAExtension(
- sources=['rbox_iou_op.cc', 'rbox_iou_op.cu'],
- extra_compile_args={'cxx': ['-DPADDLE_WITH_CUDA']}))
- else:
- setup(
- name='rbox_iou_ops',
- ext_modules=CppExtension(sources=['rbox_iou_op.cc']))
|
