TIL: How to measure memory usage from your PyTorchmodel without running it

Fantastic tip from Alban, particularly useful when you have a giant model and limited VRAM.

The short answer is this 30 lines TorchDispatchMode that tracks all Tensor memory use

https://dev-discuss.pytorch.org/t/how-to-measure-memory-usage-from-your-model-without-running-it/2024

Discover more from Ian’s Blog

Subscribe now to keep reading and get access to the full archive.

Continue reading