I don’t think there’s an easy way to do that in PyTorch today. One way you could save memory is to to use torch.utils.checkpoint, which re-computes the intermediate activations.
I think it would make sense for there to be a checkpointing option in PyTorch that saves activations on CPU to avoid using too much GPU memory. Please feel free to open us a feature request at Issues · pytorch/pytorch · GitHub.
If you do want to go down the route of manually implementing this by yourself, one (very tedious) way is to

For each operator you’re using, write a custom autograd.Function

the custom autograd.Function has a forward pass. You’d have to manually write a backward pass for it, but that can be doable by copying and pasting the relevant code from pytorch/derivatives.yaml at master · pytorch/pytorch · GitHub

When saving tensors for the backward pass, convert them to CPU before saving.