Update QwenImageTransformerBlock-fa3-kernels-custom-op/example.py

QwenImageTransformerBlock-fa3-kernels-custom-op/example.py

@@ -5,7 +5,7 @@ from diffusers import QwenImagePipeline
 _flash_attn_func = get_kernel("kernels-community/vllm-flash-attn3").flash_attn_func
 
 @torch.library.custom_op("flash::flash_attn_func", mutates_args=())
-def flash_attn_func(q, k, v):
+def flash_attn_func(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Tensor:
     outputs, lse = _flash_attn_func(q, k, v)
     return outputs