I think Microsoft released a PyTorch package some time ago, where intermediate tensors could be pushed to the CPU temporarily to reduce the GPU memory usage.
However, I can’t remember the name at the moment and don’t know if it’s still maintained.
That being said, you could trace compute for memory via torch.utils.checkpoint. 