PyTorch: Training your first Convolutional Neural Network (CNN)