Update attention_processor.py

refactor of picking a set element
huggingface · Jan 15, 2025 · 2d2411e · 2d2411e
1 parent 36d34fd
commit 2d2411e
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/diffusers/models/attention_processor.py b/src/diffusers/models/attention_processor.py
@@ -402,7 +402,7 @@ def set_use_memory_efficient_attention_xformers(
                     dtype = None
                     if attention_op is not None:
                         op_fw, op_bw = attention_op
-                        dtype = list(op_fw.SUPPORTED_DTYPES)[0]
+                        dtype, *_ = op_fw.SUPPORTED_DTYPES
                     q = torch.randn((1, 2, 40), device="cuda", dtype=dtype)
                     _ = xformers.ops.memory_efficient_attention(q, q, q)
                 except Exception as e: