Checkpointing is not compatible with .grad() or when an inputs parameter is passed to .backward()

[vrobot]

memory-efficient-attention-pytorch/memory_efficient_attention_pytorch/memory_efficient_attention.py

Line 95 in 35559a0

summarize_qkv_fn = checkpointed_summarize_qkv_chunk if needs_backwards else summarize_qkv_chunk

Should this be:
summarize_qkv_fn = summarize_qkv_chunk if needs_backwards else checkpointed_summarize_qkv_chunk
instead of:
summarize_qkv_fn = checkpointed_summarize_qkv_chunk if needs_backwards else summarize_qkv_chunk

[YuuJinnA]

Nice Job