Solved by ptrblck in post #2 
              
              
                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…