Model runs, but torchinfo will not. Possilbe issue with setting in_features in nn.Linear layer

Thanks everyone for the help! Basically, my model works, but I cannot run torchinfo. It seems there is an issue with setting the value of nn.Linear( in_features = 2560 …). I have done so manually to fix another error, and ideally, I think it would be better to assign this value dynamically. This results in a matrix multiplication error, which doesn’t make sense given that the model runs.

Here is model:

import torch
from torch import nn

class TinyVGG(nn.Module):
    """
    Model architecture copying TinyVGG from: 
    https://poloclub.github.io/cnn-explainer/
    """
    def __init__(self, input_shape: int, hidden_units: int, output_shape: int) -> None:
        super().__init__()
        self.conv_block_1 = nn.Sequential(
            nn.Conv2d(in_channels=input_shape, 
                      out_channels=hidden_units, 
                      kernel_size=3, # how big is the square that's going over the image?
                      stride=1, # default
                      padding=1), # options = "valid" (no padding) or "same" (output has same shape as input) or int for specific number 
            nn.ReLU(),
            nn.Conv2d(in_channels=hidden_units, 
                      out_channels=hidden_units,
                      kernel_size=3,
                      stride=1,
                      padding=1),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2,
                         stride=2) # default stride value is same as kernel_size
        )
        self.conv_block_2 = nn.Sequential(
            nn.Conv2d(hidden_units, hidden_units, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.Conv2d(hidden_units, hidden_units, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(2)
        )
        self.classifier = nn.Sequential(
            nn.Flatten(),
            # Where did this in_features shape come from? 
            # It's because each layer of our network compresses and changes the shape of our inputs data.
            nn.Linear(in_features= 2560, # modified from: hidden_units*16*16,
                      # above is same as 10*130*174
                      out_features=output_shape)
        )

    def forward(self, x: torch.Tensor):
        x = self.conv_block_1(x)
        # print(x.shape)
        x = self.conv_block_2(x)
        # print(x.shape)
        x = self.classifier(x)
        # print(x.shape)
        return x
        # return self.classifier(self.conv_block_2(self.conv_block_1(x))) # <- leverage the benefits of operator fusion

torch.manual_seed(42)
model_0 = TinyVGG(input_shape=128, # number of color channels (3 for RGB) - modified
                  hidden_units=10, 
                  output_shape=len(train_data_custom.classes)).to(device)
model_0 

calling torchinfo and error message:

# Install torchinfo if it's not available, import it if it is

import torchinfo

from torchinfo import summary

summary(model_0, input_size=[1, 128, 520, 696]) # do a test pass through of an example input size

{
	"name": "RuntimeError",
	"message": "Failed to run torchinfo. See above stack traces for more details. Executed layers up to: [Sequential: 1, Conv2d: 2, ReLU: 2, Conv2d: 2, ReLU: 2, MaxPool2d: 2, Sequential: 1, Conv2d: 2, ReLU: 2, Conv2d: 2, ReLU: 2, MaxPool2d: 2, Flatten: 2]",
	"stack": "---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
File ~/Library/Python/3.9/lib/python/site-packages/torchinfo/torchinfo.py:295, in forward_pass(model, x, batch_dim, cache_forward_pass, device, mode, **kwargs)
    294 if isinstance(x, (list, tuple)):
--> 295     _ = model(*x, **kwargs)
    296 elif isinstance(x, dict):

File ~/Library/Python/3.9/lib/python/site-packages/torch/nn/modules/module.py:1518, in Module._wrapped_call_impl(self, *args, **kwargs)
   1517 else:
-> 1518     return self._call_impl(*args, **kwargs)

File ~/Library/Python/3.9/lib/python/site-packages/torch/nn/modules/module.py:1568, in Module._call_impl(self, *args, **kwargs)
   1566     args = bw_hook.setup_input_hook(args)
-> 1568 result = forward_call(*args, **kwargs)
   1569 if _global_forward_hooks or self._forward_hooks:

Cell In[20], line 48, in TinyVGG.forward(self, x)
     47 # print(x.shape)
---> 48 x = self.classifier(x)
     49 # print(x.shape)

File ~/Library/Python/3.9/lib/python/site-packages/torch/nn/modules/module.py:1518, in Module._wrapped_call_impl(self, *args, **kwargs)
   1517 else:
-> 1518     return self._call_impl(*args, **kwargs)

File ~/Library/Python/3.9/lib/python/site-packages/torch/nn/modules/module.py:1568, in Module._call_impl(self, *args, **kwargs)
   1566     args = bw_hook.setup_input_hook(args)
-> 1568 result = forward_call(*args, **kwargs)
   1569 if _global_forward_hooks or self._forward_hooks:

File ~/Library/Python/3.9/lib/python/site-packages/torch/nn/modules/container.py:215, in Sequential.forward(self, input)
    214 for module in self:
--> 215     input = module(input)
    216 return input

File ~/Library/Python/3.9/lib/python/site-packages/torch/nn/modules/module.py:1518, in Module._wrapped_call_impl(self, *args, **kwargs)
   1517 else:
-> 1518     return self._call_impl(*args, **kwargs)

File ~/Library/Python/3.9/lib/python/site-packages/torch/nn/modules/module.py:1568, in Module._call_impl(self, *args, **kwargs)
   1566     args = bw_hook.setup_input_hook(args)
-> 1568 result = forward_call(*args, **kwargs)
   1569 if _global_forward_hooks or self._forward_hooks:

File ~/Library/Python/3.9/lib/python/site-packages/torch/nn/modules/linear.py:114, in Linear.forward(self, input)
    113 def forward(self, input: Tensor) -> Tensor:
--> 114     return F.linear(input, self.weight, self.bias)

RuntimeError: mat1 and mat2 shapes cannot be multiplied (1x226200 and 2560x3)

The above exception was the direct cause of the following exception:

RuntimeError                              Traceback (most recent call last)
Cell In[22], line 4
      2 import torchinfo
      3 from torchinfo import summary
----> 4 summary(model_0, input_size=[1, 128, 520, 696]) # do a test pass through of an example input size 

File ~/Library/Python/3.9/lib/python/site-packages/torchinfo/torchinfo.py:223, in summary(model, input_size, input_data, batch_dim, cache_forward_pass, col_names, col_width, depth, device, dtypes, mode, row_settings, verbose, **kwargs)
    216 validate_user_params(
    217     input_data, input_size, columns, col_width, device, dtypes, verbose
    218 )
    220 x, correct_input_size = process_input(
    221     input_data, input_size, batch_dim, device, dtypes
    222 )
--> 223 summary_list = forward_pass(
    224     model, x, batch_dim, cache_forward_pass, device, model_mode, **kwargs
    225 )
    226 formatting = FormattingOptions(depth, verbose, columns, col_width, rows)
    227 results = ModelStatistics(
    228     summary_list, correct_input_size, get_total_memory_used(x), formatting
    229 )

File ~/Library/Python/3.9/lib/python/site-packages/torchinfo/torchinfo.py:304, in forward_pass(model, x, batch_dim, cache_forward_pass, device, mode, **kwargs)
    302 except Exception as e:
    303     executed_layers = [layer for layer in summary_list if layer.executed]
--> 304     raise RuntimeError(
    305         \"Failed to run torchinfo. See above stack traces for more details. \"
    306         f\"Executed layers up to: {executed_layers}\"
    307     ) from e
    308 finally:
    309     if hooks:

RuntimeError: Failed to run torchinfo. See above stack traces for more details. Executed layers up to: [Sequential: 1, Conv2d: 2, ReLU: 2, Conv2d: 2, ReLU: 2, MaxPool2d: 2, Sequential: 1, Conv2d: 2, ReLU: 2, Conv2d: 2, ReLU: 2, MaxPool2d: 2, Flatten: 2]"
}

The issue is unrelated to torchinfo as the model execution itself already fails with a shape mismatch:

x = torch.randn(1, 128, 520, 696)
out = model_0(x)
# RuntimeError: mat1 and mat2 shapes cannot be multiplied (1x226200 and 2560x1)

Ah, yes you are right!
I was passing in the wrong dimensions.
Thanks!