Image Classification with custom Dataset

I think you copied that network code from the original Pytorch tutorial https://pytorch.org/tutorials/beginner/blitz/neural_networks_tutorial.html#sphx-glr-beginner-blitz-neural-networks-tutorial-py . In that site the image is 3 * 32 * 32 but yours is 3 * 200* * 200, so when transiting to the first fully connected layer you have to flatten the last output of you CNN layer. This might help https://discuss.pytorch.org/t/about-nn-linear-16-5-5-120/62206/3