I have this problem and cannot solve it.
This is my Dataset:
class Data(Dataset):
    def __init__(self, df, data_dir, transform):
        super().__init__()
        self.df = df
        self.data_dir = data_dir
        self.transform = transform

    def __len__(self):
        return len(self.df)

    def __getitem__(self, index):
        img_name = self.df.id[index]
        label = self.df.broken[index]

        img_path = os.path.join(str(self.data_dir), str(img_name) + '.png')
        img = mpimg.imread(img_path)
        img = self.transform(img)
        return img, label

This is my model:
class SkullCNN(nn.Module):
    def __init__(self):
        super(SkullCNN, self).__init__()
        self.conv1 = nn.Sequential(
            nn.Conv2d(3, 32, 4, 2, 0),
            nn.BatchNorm2d(32),
            nn.ReLU()
        )
        # 1 + (32 - 4 + 0)/2 = 15
        # 32 * 15 * 15
        self.conv2 = nn.Sequential(
            nn.Conv2d(32, 64, 3, 2, 0),
            nn.BatchNorm2d(64),
            nn.ReLU()
        )
        # 1 + (15 - 3 + 0)/2 = 7
        # 64 * 7 * 7
        self.conv3 = nn.Sequential(
            nn.Conv2d(64, 128, 3, 2, 0),
            nn.BatchNorm2d(128),
            nn.ReLU()
        )
        # 1 + (7 - 3 + 0)/2 = 3
        # 128 * 3 * 3
        self.conv4 = nn.Sequential(
            nn.Conv2d(128, 256, 3, 2, 0),
            nn.BatchNorm2d(256),
            nn.ReLU()
        )
        #  1 + (3 - 3 + 0)/2 = 1
        # 256 * 1 * 1
        
        self.fc = nn.Sequential(
            nn.Linear(256*1*1, 1024),
            nn.ReLU(),
            nn.Dropout(p=0.2),
            nn.Linear(1024,2)
        )
    def forward(self, x):
        x = self.conv1(x)

        x = self.conv2(x)

        x = self.conv3(x)

        x = self.conv4(x)

        x = x.view(x.shape[0],-1)
        x = self.fc(x)
        return x

This is the code block where the error is generated:
num_epochs = 25
num_classes = 2
batch_size = 128

for epoch in range(num_epochs):
    for i, (images, labels) in enumerate(train_loader):
        images = images.to(device)
        labels = labels.to(device)
#         print(images[0].shape)
        
        out = model(images)
        loss = criterion(out, labels)
        
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
    print('Epoch: {}/{}, Loss: {}'.format(epoch+1, num_epochs, loss.item()))

The error output is:
---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
<ipython-input-72-15d7d37170f1> in <module>()
      9 #         print(images[0].shape)
     10 
---> 11         out = model(images)
     12         loss = criterion(out, labels)
     13 

6 frames
/usr/local/lib/python3.6/dist-packages/torch/nn/modules/conv.py in _conv_forward(self, input, weight)
    348                             _pair(0), self.dilation, self.groups)
    349         return F.conv2d(input, weight, self.bias, self.stride,
--> 350                         self.padding, self.dilation, self.groups)
    351 
    352     def forward(self, input):

RuntimeError: Given groups=1, weight of size [32, 3, 4, 4], expected input[64, 1, 340, 340] to have 3 channels, but got 1 channels instead

Please help me, I’m learning PyTorch. Thank you so much, also I see some similar problems here, but I can’t adapt it to my code.