GPU memory that model uses

To calculate the memory requirement for all parameters and buffers, you could simply sum the number of these and multiply by the element size:

mem_params = sum([param.nelement()*param.element_size() for param in model.parameters()])
mem_bufs = sum([buf.nelement()*buf.element_size() for buf in model.buffers()])
mem = mem_params + mem_bufs # in bytes

However, this will not include the peak memory usage for the forward and backward pass (if that’s what you are looking for).

7 Likes