Solved by ptrblck in post #4 
              
              
                Thanks for the shape information. 
The weight matrices alone will take approx. 165GB. 
If we ignore the bias, you can calculate the memory footprint of e.g. the last linear layer’s weight matrix using: 
256*22*22 * 256*32*32 * 4 / 1024**3 = 121 GB

I assume you are using float32 values, which explai…