Update aten/src/ATen/native/transformers/cuda/attention_backward.cu

eqy · seemethere · pytorchmergebot · commit b53102bf5c54 · 2024-04-26T22:28:36.000Z
Co-authored-by: Eli Uriegas &lt;1700823+seemethere@users.noreply.github.com&gt;
diff --git a/aten/src/ATen/native/transformers/cuda/attention_backward.cu b/aten/src/ATen/native/transformers/cuda/attention_backward.cu
@@ -43,7 +43,11 @@
 #include <ATen/native/transformers/cuda/mem_eff_attention/pytorch_utils.h>
 #endif
 
+#ifdef __HIP_PLATFORM_AMD__
+#include <ATen/native/cudnn/hip/MHA.h>
+#else
 #include <ATen/native/cudnn/MHA.h>
+#endif
 
 namespace at::native {
 
