ptrblck:

I would thus recommend to perform an example training step with the shapes you are planning to use and check the memory usage e.g. via torch.cuda.memory_summary()


The problem is usually these models I am referring to are too big to fit in memory, and the point of calculating the memory is we are deciding how to distribute each sub-model onto multiple GPUs. So, I can’t load the entire model first to calculate the size.
Do you have any code for the first method of calculating memory?