check tensors on accelerator

lauri9 · lauri9 · commit 0dce2cb3ac62 · 2025-10-24T07:50:58.000-05:00
diff --git a/src/diffusers/models/attention_dispatch.py b/src/diffusers/models/attention_dispatch.py
@@ -1418,7 +1418,7 @@ def _flash_varlen_attention_3(
 
 @_AttentionBackendRegistry.register(
     AttentionBackendName.AITER,
-    constraints=[_check_device, _check_qkv_dtype_bf16_or_fp16, _check_shape],
+    constraints=[_check_device_cuda, _check_qkv_dtype_bf16_or_fp16, _check_shape],
 )
 def _aiter_flash_attention(
     query: torch.Tensor,

Original file line number	Diff line number	Diff line change
`@@ -1418,7 +1418,7 @@ def _flash_varlen_attention_3(`
`1418`	`1418`
`1419`	`1419`	`@_AttentionBackendRegistry.register(`
`1420`	`1420`	`AttentionBackendName.AITER,`
`1421`		`- constraints=[_check_device, _check_qkv_dtype_bf16_or_fp16, _check_shape],`
	`1421`	`+ constraints=[_check_device_cuda, _check_qkv_dtype_bf16_or_fp16, _check_shape],`
`1422`	`1422`	`)`
`1423`	`1423`	`def _aiter_flash_attention(`
`1424`	`1424`	`query: torch.Tensor,`