RuntimeError: Given groups=1, weight of size [6, 1, 5, 5], expected input[128, 3, 218, 178] to have 1 channels, but got 3 channels instead

[ptrblck]

Yes, that was expected and matches my second point:

@ptrblck

and fix the in_features of self.fc1 in case you are running into a shape mismatch afterwards.

Change self.fc1 to:

self.fc1 = nn.Linear(34944, 120)