A few things to be careful about when using this are:

At least one input to checkpoint function must require gradients. You can make your model input requires grad for example if the parameters are “hidden” inside the model.
It will allow you to save the memory from the intermediary results inside your function (but not the ones that are inputs or outputs). So to get any benefit, your function must be big enough.
It will recompute these intermediary results during the backward and only keep them around while the backward for the checkpoint is computed. So you cannot put all your function inside a single checkpoint otherwise the max memory requirement will be the same as before.

An example is if you use resnet network that contains many blocks of conv + conv + relu + bn. Then you can checkpoint each of these blocks to reduce the memory usage without a big speed penalty.