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.