Memory estimates for recursive models are inconsistent with actual GPU usage

[Ethan-yt]

Describe the bug Memory estimates are inconsistent with actual GPU usage for recursive models

To Reproduce Example code:

import torch
from torch.nn import Sequential, Linear

class MyModel(torch.nn.Module):
    def __init__(self, n):
        super().__init__()
        self.linear = Linear(1024, 1024)

    def forward(self, x):
        for i in range(100):
            x = self.linear(x)
        return x

model = MyModel().cuda()

a = model(torch.zeros((10240, 1024)).cuda())
print(torch.cuda.memory_allocated()/1024/1024)

from torchinfo import summary

summary(model,input_size=(10240, 1024))

torch reported: 4044.00390625

Input size (MB): 41.94 Forward/backward pass size (MB): 83.89 Params size (MB): 4.20 Estimated Total Size (MB): 130.03

If i change 100 to 1 all looks good:

torch reported: 92.01171875

Input size (MB): 41.94 Forward/backward pass size (MB): 83.89 Params size (MB): 4.20 Estimated Total Size (MB): 130.03

[TylerYep]

Good observation. I would appreciate any PRs addressing this issue!

[Ethan-yt]

looks like it neglects the intermediate tensors stored for backward propagation Estimated: Forward/backward pass size (MB): 83.89 = 102401024(input shape)4(byte)*2(grad)=83,886,080

torch reported: 4044.00390625MB = 4.003MB(W and b) + 40MB(input) + 40MB(intermediate)* 100(iterations)

[TylerYep]

This issue should be resolved in #181, but I'll wait for the release and a verification before closing