Note at Section 03, using a trick for input and output

[aronvandepol]

Hi all,

With the pytorch version 1.11 (and 1.12), the trick that Daniel uses (hidden_units*7*7) doesn;'t work. It worked I believe because the output in 1.10 of Conv_layer_2 =[1,10,7,7]. Multiplying each unit 10*7*7 = 490 and delivers [1,490] and thus solving this by using hidden_units*7*7 works in 1.10.

In 1.11 and 1.12, the output of conv_layer_2 is however is [10, 7, 7], leading to 7*7 and a size of [10*49]. Hence, you cannot solve the input by doing hidden*7*7 (results in 490) but rather, simply 7*7.

thus the linear layer becomes:

nn.Linear(in_features=7*7, out_features=output_shape)

Using this the shapes match and it will work on a single image,

Yet when training you will need the hidden*7*7 setup as it wont work otherwise.

[mrdbourke]

Thank you for this!

I didn't realise.

I'm going to add this as an issue and inspect it shortly.

See the issue here: #71

[shisirkha]

@aronvandepol Maybe this is not that needed i have a solution to it. (NOTE: this solution is only tested on food101 dataset mini)

We can just multiply both input and output shape of the model in an example (This totally works):

from torch import nn
import torch
#Let us inherite from nn.Module
class TinyVGG(nn.Module):
  def __init__(self,input_shape: int,hidden_units: int,output_shape: int):
    super().__init__()
    self.layer_1 = nn.Sequential(
        nn.Conv2d(input_shape,hidden_units,3,1,1),
        nn.ReLU(),
        nn.Conv2d(hidden_units,hidden_units,3,1,1),
        nn.ReLU(),
        nn.MaxPool2d(2,2)
    )
    self.layer_2 = nn.Sequential(
        nn.Conv2d(hidden_units,hidden_units,3,1,1),
        nn.ReLU(),
        nn.Conv2d(hidden_units,hidden_units,3,1,1),
        nn.ReLU(),
        nn.MaxPool2d(2,2)
    )
    self.classifier = nn.Sequential(
        nn.Flatten(),
        nn.Linear(hidden_units*16*16, output_shape) #<--- here is the change we multiply both hidden shape after Flattening it. this might fix the problem
    )
  def forward(self,x):
    x = self.layer_1(x)
    x = self.layer_2(x)
    x = self.classifier(x)
    return x

@mrdbourke Does this method solve the problem?

[AlexeusPodbeltsev]

Yo! Did you unsqueeze an input tensor? Cause I had same problem as you have explained and found a solution similar to yours, but when i ran it, i found that the output of the model is 10x10 matrix, not a vector. So after a few more attempts to understand where I went wrong, i finally got an answer. so keep in mind the batch_num in an input tensor.
model_2(rand_image_tensor.unsqueeze(0).to(device))

My PyTorch version is 2.5.1
I know that this discussion is quite old, but maybe it will be helpfull for a new ones))