a small typo and fix

[liguohao96]

I compiled from source and run a simple test but failed at this line

flash-attention/flash_attn/flash_attn_interface.py

Line 853 in 0dfb281

return out if not return_softmax else (out, softmax_lse, S_dmask)

with the following error: unboundLocalError: cannot access local variable 'out' where it is not associated with a value

I believe it is caused by the wrong ident of this line

flash-attention/flash_attn/flash_attn_interface.py

Line 852 in 0dfb281

out = out_padded[..., :head_size_og]

we should move this line out of if block.

[tridao]

I see, thanks for catching it. Can you send a PR?

[liguohao96]

Here is the PR #1392 .
BTW, I tried to add some 'forward-only' tests, but finally I gave up because I failed to find an elegant way to reuse current test functions.

[jiagaoxiang]

Could you please merge this PR ASAP to unblock? Thank you very much!