Change how we detect device

pytorch-labs · Nov 29, 2024 · e4985e2 · e4985e2
1 parent e8e20f6
commit e4985e2
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/tritonbench/kernels/triton_fused_attention_vanilla.py b/tritonbench/kernels/triton_fused_attention_vanilla.py
@@ -4,7 +4,9 @@
 
 
 def is_hip():
-    return triton.runtime.driver.active.get_current_target().backend == "hip"
+    # use pytorch to detect current device
+    # return triton.runtime.driver.active.get_current_target().backend == "hip"
+    return bool(torch.version.hip)
 
 
 @triton.jit