Hi @chuanyangjin ,

When I run your repo with gradient checkpointing, I find that it is super slow with multiple GPUs.
I think the culprit is torch.checkpoint. Is there an alternative?