Incorrect code example in beginner tutorial (Build the Neural Network)

📚 Documentation

Hey, I was following along the beginner tutorial and attempted to run the following code sample here:
https://pytorch.org/tutorials/beginner/basics/buildmodel_tutorial.html

X = torch.rand(1, 28, 28, device=device)
logits = model(X)
pred_probab = nn.Softmax(dim=1)(logits)
y_pred = pred_probab.argmax(1)
print(f"Predicted class: {y_pred}")

Running this example gave me the following RuntimeError

RuntimeError: The size of tensor a (28) must match the size of tensor b (784) at non-singleton dimension 2

After looking up the torch.rand docs https://pytorch.org/docs/stable/generated/torch.rand.html I attempted to change the parameters to something that seemed more correct considering the input tensor to the example NeuralNetwork class.

X = torch.rand(1, 28*28, device=device)

This seems to let the code example execute as expected (by the docs).

I will note however I am a TOTAL beginner here so if this isn't the right solution then please let me know!

Start with the code sample in the linked Build the Neural Network beginner tutorial and reproduce the reported RuntimeError. Compare the sample input shape with the NeuralNetwork example and update the tutorial so the example runs successfully without the tensor-size error.