Hi Vimal!



 vimal4353:

   Linear(in_features=784, out_features=128, bias=True)



The first layer of your model is a Linear with in_features = 784.  It.
therefore expects to be passed a batch of 1-d vectors of length 784.
(More precisely, a Linear expects a tensor whose last dimension is
in_features and that has any number of leading “batch” dimensions,
including zero.  So you can think of this as being a potentially
multidimensional “batch” of 1-d vectors, if you will.)
It looks like you are passing in a image with width = height = 28.
You need to flatten() your 28x28 image into a length-784 vector (or,
more generally, flatten() the last two dimensions of your batch of
images) before passing it to your Linear.
Note, you can include the flattening operation as a “layer” in your
model by using the Module (class or function-object) version, Flatten.
Consider:
>>> import torch
>>> torch.__version__
'1.9.0'
>>> batch_of_3_images = torch.randn (3, 28, 28)
>>> batch_of_3_images.shape
torch.Size([3, 28, 28])
>>> lin = torch.nn.Linear (in_features = 784, out_features = 128, bias = True)
>>> flat = torch.nn.Flatten (start_dim = -2)
>>> flat_batch = flat (batch_of_3_images)
>>> flat_batch.shape
torch.Size([3, 784])
>>> lin (flat_batch).shape
torch.Size([3, 128])

Best.
K. Frank