Exporting Stable diffusion 's VAE model Failed

[TSCCompiler]

I'm using torch 2.1.0.dev20230425+cpu and diffuser 0.16  to build stable diffusion v1_5, But I got the following error:
assert ( AssertionError: Unsupported function type scaled_dot_product_attention

I print symbolic traced graph , and I found vae module is using torch.c.scaled_dot_procuct_attention op

scaled_dot_product_attention = torch._C._nn.scaled_dot_product_attention(permute, permute_1, permute_2, dropout_p = 0.0, is_causal = False); permute_1 = permute_2 = None
I can attach the code which can make this happend
output.zip

[TSCCompiler]

I have figure this out, I made diffusers use torch.baddbmm instead of scaled_dot_product_attention

[ford442]

would that work for the webgpu compile as well?
i have the same error on CPU or CUDA..