No errors before saving, but an error during testing

When training my model, I didn’t encounter any errors. However, after saving and loading the model, I encountered an error when testing it with the loaded model. mat1 and mat2 shapes cannot be multiplied (10x576 and 5760x12)

Could you post your model definition as well as the input shapes here, please?

The input size of the model = torch.Size([3, 224, 224])

class TinyVGG(nn.Module):
    def __init__(self, input_shape: int, hidden_units: int, output_shape: int):
        super().__init__()
        self.block_1 = nn.Sequential(
            nn.Conv2d(in_channels=input_shape, out_channels=hidden_units, kernel_size=3, stride=1, padding=0),
            nn.ReLU(),
            nn.Conv2d(in_channels=hidden_units, out_channels=hidden_units, kernel_size=3, stride=1, padding=0),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2)
        )
        self.block_2 = nn.Sequential(
            nn.Conv2d(in_channels=hidden_units, out_channels=hidden_units, kernel_size=3, stride=1, padding=0),
            nn.ReLU(),
            nn.Conv2d(in_channels=hidden_units, out_channels=hidden_units, kernel_size=3, stride=1, padding=0),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2)
        )
        self.classifier = nn.Sequential(
            nn.Flatten(),
            nn.Linear(in_features=hidden_units*24*24, out_features=output_shape)
        )

    def forward(self, x: torch.Tensor):
        return self.classifier(self.block_2(self.block_2(self.block_1(x))))


model_0 = TinyVGG(input_shape=3, hidden_units=10, output_shape=len(class_names))

Thanks for the code!
I cannot reproduce any issues and the model works for me after loading it and calling .eval() on it (assuming that’s also your workflow):

import torch
import torch.nn as nn

class TinyVGG(nn.Module):
    def __init__(self, input_shape: int, hidden_units: int, output_shape: int):
        super().__init__()
        self.block_1 = nn.Sequential(
            nn.Conv2d(in_channels=input_shape, out_channels=hidden_units, kernel_size=3, stride=1, padding=0),
            nn.ReLU(),
            nn.Conv2d(in_channels=hidden_units, out_channels=hidden_units, kernel_size=3, stride=1, padding=0),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2)
        )
        self.block_2 = nn.Sequential(
            nn.Conv2d(in_channels=hidden_units, out_channels=hidden_units, kernel_size=3, stride=1, padding=0),
            nn.ReLU(),
            nn.Conv2d(in_channels=hidden_units, out_channels=hidden_units, kernel_size=3, stride=1, padding=0),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2)
        )
        self.classifier = nn.Sequential(
            nn.Flatten(),
            nn.Linear(in_features=hidden_units*24*24, out_features=output_shape)
        )

    def forward(self, x: torch.Tensor):
        return self.classifier(self.block_2(self.block_2(self.block_1(x))))


model_0 = TinyVGG(input_shape=3, hidden_units=10, output_shape=10)

# test before saving
x = torch.randn(1, 3, 224, 224)
out = model_0(x)
out.mean().backward()

# save model
torch.save(model_0.state_dict(), "tmp.pt")

# restore model
model = TinyVGG(input_shape=3, hidden_units=10, output_shape=10)
model.load_state_dict(torch.load("tmp.pt"))
model.eval()

# test after restoring
x = torch.randn(1, 3, 224, 224)
out = model(x)
out.mean().backward()
with torch.inference_mode():
    pred = model_0(data)
pred

after loading this code i and testing with an data
the error is

RuntimeError Traceback (most recent call last)
Cell In[88], line 2
1 with torch.inference_mode():
----> 2 pred = model_0(data)
3 pred

File ~/anaconda3/lib/python3.12/site-packages/torch/nn/modules/module.py:1553, in Module._wrapped_call_impl(self, *args, **kwargs)
1551 return self._compiled_call_impl(*args, **kwargs) # type: ignore[misc]
1552 else:
→ 1553 return self._call_impl(*args, **kwargs)

File ~/anaconda3/lib/python3.12/site-packages/torch/nn/modules/module.py:1562, in Module._call_impl(self, *args, **kwargs)
1557 # If we don’t have any hooks, we want to skip the rest of the logic in
1558 # this function, and just call forward.
1559 if not (self._backward_hooks or self._backward_pre_hooks or self._forward_hooks or self._forward_pre_hooks
1560 or _global_backward_pre_hooks or _global_backward_hooks
1561 or _global_forward_hooks or _global_forward_pre_hooks):
→ 1562 return forward_call(*args, **kwargs)
1564 try:
1565 result = None

Cell In[65], line 27, in TinyVGG.forward(self, x)
26 def forward(self, x: torch.Tensor):
—> 27 return self.classifier(self.block_2(self.block_2(self.block_1(x))))

File ~/anaconda3/lib/python3.12/site-packages/torch/nn/modules/module.py:1553, in Module._wrapped_call_impl(self, *args, **kwargs)
1551 return self._compiled_call_impl(*args, **kwargs) # type: ignore[misc]
1552 else:
→ 1553 return self._call_impl(*args, **kwargs)

File ~/anaconda3/lib/python3.12/site-packages/torch/nn/modules/module.py:1562, in Module._call_impl(self, *args, **kwargs)
1557 # If we don’t have any hooks, we want to skip the rest of the logic in
1558 # this function, and just call forward.
1559 if not (self._backward_hooks or self._backward_pre_hooks or self._forward_hooks or self._forward_pre_hooks
1560 or _global_backward_pre_hooks or _global_backward_hooks
1561 or _global_forward_hooks or _global_forward_pre_hooks):
→ 1562 return forward_call(*args, **kwargs)
1564 try:
1565 result = None

File ~/anaconda3/lib/python3.12/site-packages/torch/nn/modules/container.py:219, in Sequential.forward(self, input)
217 def forward(self, input):
218 for module in self:
→ 219 input = module(input)
220 return input

File ~/anaconda3/lib/python3.12/site-packages/torch/nn/modules/module.py:1553, in Module._wrapped_call_impl(self, *args, **kwargs)
1551 return self._compiled_call_impl(*args, **kwargs) # type: ignore[misc]
1552 else:
→ 1553 return self._call_impl(*args, **kwargs)

File ~/anaconda3/lib/python3.12/site-packages/torch/nn/modules/module.py:1562, in Module._call_impl(self, *args, **kwargs)
1557 # If we don’t have any hooks, we want to skip the rest of the logic in
1558 # this function, and just call forward.
1559 if not (self._backward_hooks or self._backward_pre_hooks or self._forward_hooks or self._forward_pre_hooks
1560 or _global_backward_pre_hooks or _global_backward_hooks
1561 or _global_forward_hooks or _global_forward_pre_hooks):
→ 1562 return forward_call(*args, **kwargs)
1564 try:
1565 result = None

File ~/anaconda3/lib/python3.12/site-packages/torch/nn/modules/linear.py:117, in Linear.forward(self, input)
116 def forward(self, input: Tensor) → Tensor:
→ 117 return F.linear(input, self.weight, self.bias)

RuntimeError: mat1 and mat2 shapes cannot be multiplied (10x576 and 5760x12)

No errors using my code:

with torch.inference_mode():
    out = model(x)
print(out.shape)
# torch.Size([1, 10])

with torch.inference_mode():
    out = model_0(x)
print(out.shape)
# torch.Size([1, 10])

Here are the links to the GitHub repository and the Kaggle dataset.

I am implementing the same thing in the above code, but I am getting this error.

Thank you, sir. I solved the problem. The error was because when testing with test data, I didn’t provide the unsquace. Thank you for helping me with this

with torch.inference_mode():
    pred = model(data.unsqueeze(0))
pred