My question concerns the logic of building a neural network.

[Duwimpuhwe]

With created instance of built model, the input shape is equivalent to the size number of pixels of a flatten image 28*28=784, but during the training and prediction the batch of 32 images is flattened and passed to the model (through the first layer onwards). how does the batch of image processed while the create models has the first layer with neurons matching with the size of the single images?

[mrdbourke]

Hi @Duwimpuhwe ,

Good question!

That can be a bit confusing when you start to work with batches.

But the way torch.nn.Linear (and many other layers) are designed, they can work with multiple sized inputs on the first dimension (most often the batch dimension).

E.g. (32, 784) could be a tensor of batch size 32 of flattened images with size 784.

See the torch.nn.Linear documentation screenshot for more:

I'd also try playing around with some example code/tensors to get a feel of what the inputs/outputs of the layers feel like.

For example:

import torch
from torch import nn

m = nn.Linear(20, 30)
input = torch.randn(128, 20)
output = m(input)
print(output.size())