Based on the error message, it looks like output and label do not have the same batch size.
Could you print the shapes of train, labels, and output and check, why the shapes are different?
I couldn’t find any obvious error in reshaping by skimming through your code.
PS: I’m not sure, if it’s on purpose, but you are not using any activation function on the output of self.fc1.