gradient checkpointing with timm

[DonkeyShot21]

looks very easy, we can just call this function:
https://github.com/rwightman/pytorch-image-models/blob/8ff45e41f7a6aba4d5fdadee7dc3b7f2733df045/timm/models/resnet.py#L697

[vturrisi]

Bumping this up as we plan to do it next week.