I looked into torch.utils.checkpoint, it does satisfy what I needed.
I also tried searching again with different keywords and found this repository https://github.com/IBM/pytorch-large-model-support 45 It enables memory swapping between CPU memory and GPU (similar to memory swap between CPU RAM and storage memory). I will leave the link here maybe it can help someone later on.
I will look into microsoft releases, if I happen to find anything related to the topic I’ll try to remember to update the post.
Thanks for the help