Checkpointing flag of Attention Block is not correctly set up

[jhliu17]

Checkpointing flag of Attention Block is not correctly set up

The Attention Block was always using checkpointing by setting the flag to True. This pull request fixes this issue by respecting the use_checkpoint parameter.

Before submitting

• Did you make sure title is self-explanatory and the description concisely explains the PR?
• Did you make sure your PR does only one thing, instead of bundling different changes together?
• Did you list all the breaking changes introduced by this pull request?
• Did you test your PR locally with pytest command?
• Did you run pre-commit hooks with pre-commit run -a command?

Did you have fun?

Make sure you had fun coding 🙃

[atong01]

Thanks for catching this! LGTM.