Hello,
you have to make sure that you are consistent in terms of shapes. I think the problem is within your flattening: When doing
  x = x.view(-1, self.kernel_size * self.kernel_size * 16)

the output of your last conv layer actually has to be kernel_size x kernel_size (in your case 5) with 16 channels. If this is not the case, it will shift the remaining values to the batch dimension (because of -1). But this depends on your input size which is not specified in your code.
If I’m not mistaken, as a quick workaround you could simply add a factor of 4 to the relevant lines of code (because of the mismatch 16 != 4) in your first linear layer and
 x = x.view(-1, 4 * self.kernel_size * self.kernel_size * 16)

which ensures that the batch dimension of the ground truth and predicted labels are equal. (Neglecting that this is not the prettiest solution, because it’s not generic in terms of your input image size).
Best regards