TylerYep / torchinfo

View model summaries in PyTorch!
MIT License
2.56k stars 119 forks source link

Pytorch 1.12 - torchvision - NonDynamicallyQuantizableLinear - No output produced #164

Open ArulselvanMadhavan opened 2 years ago

ArulselvanMadhavan commented 2 years ago

Describe the bug Pytorch 1.12 - torchvision - NonDynamicallyQuantizableLinear Not producing input/output shapes for this layer. I have pasted sample code and sample output below

class MultiheadAttentionWrapper(torch.nn.Module):
    def __init__(self, layer):
        super().__init__()
        self.layer = layer

    def forward(self, inputs):
        output = self.layer(inputs, inputs, inputs)
        return output[0]

    def __getitem__(self, index):
        return self.layer

model = MultiheadAttentionWrapper(torch.nn.MultiheadAttention(768,12,batch_first=True))

col_names = (
                "input_size",
                "output_size",
                "num_params",
            )
orig_results = torchinfo.summary(model, input_data=[input_tensors], verbose=0, device=device, col_names = col_names)
print(orig_results)

Expected behavior A clear and concise description of what you expected to happen.

[=============================================================================================================================
Layer (type:depth-idx)                             Input Shape               Output Shape              Param #
=============================================================================================================================
MultiheadAttentionWrapper                          [2, 197, 768]             [2, 197, 768]             --
├─MultiheadAttention: 1-1                          [2, 197, 768]             [2, 197, 768]             1,771,776
│    └─NonDynamicallyQuantizableLinear: 2-1        --                        --                        590,592
=============================================================================================================================
Total params: 2,362,368
Trainable params: 2,362,368
Non-trainable params: 0
Total mult-adds (M): 0.00
=============================================================================================================================
Input size (MB): 1.21
Forward/backward pass size (MB): 0.00
Params size (MB): 2.36
Estimated Total Size (MB): 3.57
=============================================================================================================================]

Desktop (please complete the following information):

Additional context Add any other context about the problem here.

TylerYep commented 2 years ago

What are the expected input/output values?