torch compiler not working with sage attention 2.2.0
I have pytorch 2.7.0 with cuda 12.8, python 3.12.0 under ubuntu 25.04 with rtx 5090. triton 3.3.0 from here
Any idea why this happens? I used the https://huggingface.co/Kijai/PrecompiledWheels/blob/main/sageattention-2.2.0-cp312-cp312-linux_x86_64.whl
torch.* op returned non-Tensor int call_function <function device_count at 0x7f6a4d3be8e0>
from user code:
File "/home/rtl-6/comfy/ldm/wan/model.py", line 221, in forward
y = self.self_attn(
File "/home/rtl-6/comfy/ldm/wan/model.py", line 72, in forward
x = optimized_attention(
File "/home/rtl-6/comfy/ldm/modules/attention.py", line 496, in attention_sage
out = sageattn(q, k, v, attn_mask=mask, is_causal=False, tensor_layout=tensor_layout)
File "/home/rtl-6/Python-3.12.0/comfy-env-3.12/lib/python3.12/site-packages/sageattention/core.py", line 140, in sageattn
arch = get_cuda_arch_versions()[q.device.index]
File "/home/rtl-6/Python-3.12.0/comfy-env-3.12/lib/python3.12/site-packages/sageattention/core.py", line 71, in get_cuda_arch_versions
for i in range(torch.cuda.device_count()):
Set TORCHDYNAMO_VERBOSE=1 for the internal stack trace (please do this especially if you're reporting a bug to PyTorch). For even more developer context, set TORCH_LOGS="+dynamo"