The problem is that your batch size is 1. Make it greater than 1
The calculation you made is right. So change self.fc1 = nn.Linear(in_features=3136, out_features=256) to self.fc1 = nn.Linear(in_features=784, out_features=256).
Refer to Error: Expected more than 1 value per channel when training to know why the batch size has to be greater than 1.