Format flash_attention with black (#2444)
Summary: Pull Request resolved: https://github.com/pytorch/benchmark/pull/2444
Reviewed By: adamomainz
Differential Revision: D61924710
Pulled By: bertmaher
fbshipit-source-id: 1d7ea4aad48d34890d1a4cc613a1f293fb781c4d
