[bug] attention benchmark script not work

[MaxwellDing]

I seem to be having some trouble getting the benchmark script to run. I noticed that there have been some API changes, and it looks like the benchmark script hasn't been updated to accommodate them yet. After I adjusted the API calls in the script, it ran, but I'm not seeing any results. Could you please take a look and help me figure out what's going on?

Thanks so much for your time!

https://github.com/NVIDIA/TransformerEngine/blame/6f4310d700f7445fd12d524c645b1e72fb8886f7/benchmarks/attention/benchmark_attention.py#L200C17-L200C17

output after fixed:

       cuDNN fwd+bwd (ms)  flash-attn fwd+bwd (ms)  cuDNN vs flash speedup
test_0              0.0000                        0                       0
test_1              0.0000                        0                       0
test_2              3.0056                        0                       0
test_3              0.0000                        0                       0