diff --git a/op_builder/builder.py b/op_builder/builder.py index 3b1730567eef..79692ce05878 100644 --- a/op_builder/builder.py +++ b/op_builder/builder.py @@ -616,6 +616,9 @@ def builder(self): if not self.build_for_cpu and self.enable_bf16: compile_args['cxx'].append("-DBF16_AVAILABLE") + if self.is_rocm_pytorch(): + compile_args['cxx'].append("-D__HIP_PLATFORM_AMD__=1") + cuda_ext = ExtensionBuilder(name=self.absolute_name(), sources=self.strip_empty_entries(self.sources()), include_dirs=self.strip_empty_entries(self.include_paths()),