CNN dimension error
Hi there!
I am trying to build a very basic CNN for a binary classification task however I am getting an odd dimensionality issue.
My CNN is:
class convNet(nn.Module):
    def __init__(self):
        
        super().__init__()
        self.conv2d_1 = nn.Conv2d(193, 193, kernel_size=3)
        self.conv2d_2 = nn.Conv2d(193, 193, kernel_size=3)
        self.conv2d_3 = nn.Conv2d(193, 193, kernel_size=3)
        self.maxpool = nn.MaxPool2d(2)
    
    def forward(self, x):
        x = self.conv2d_1(x)
        x = F.relu(x)
        x = self.maxpool(x)
        x = self.conv2d_2(x)
        x = F.relu(x)
        x = self.maxpool(x)
        x = self.conv2d_3(x)
        return F.relu(x)
And my lightning module is:
class classifier(pl.LightningModule):
    
    def __init__(self, learning_rate = float):
        super().__init__()
        
        self.learning_rate = learning_rate
        self.cnn =covNet()
        self.flat = nn.Flatten()
        
        self.fc1 = nn.Linear(450076, 100)
        self.fc2 = nn.Linear(100, 10)
        self.fc3 = nn.Linear(10, 1)
        
        self.dropout = nn.Dropout(p = 0.2)
        
        self.criterion = nn.BCEWithLogitsLoss()
        
        self.accuracy = tm.Accuracy()

    def prepare_batch(self, batch):
        img = batch['image'][tio.DATA]
        img = torch.squeeze(img)
        diagnosis = batch['diagnosis']
        return img, diagnosis

    def forward(self, x):
        cnn_out = self.cnn(x)
        flat = self.flat(cnn_out)
        fc1_out = self.dropout(F.relu(self.fc1(flat)))
        fc2_out = self.dropout(F.relu(self.fc2(fc1_out)))
        fc3_out = F.relu(self.fc3(fc2_out))
        return fc3_out
    
    def training_step(self, batch, batch_idx):
        x, y = self.prepare_batch(batch)
        y = y.view(y.size(0), -1)
        y = y.type(torch.float)
        y_hat = self.forward(x)
        train_loss = self.criterion(y_hat, y)
        self.log('train_loss', train_loss, prog_bar = True)
        return train_loss
    
    def validation_step(self, batch, batch_idx):
        x, y = self.prepare_batch(batch)
        y = y.view(y.size(0), -1)
        y = y.type(torch.float)
        y_hat = self.forward(x)
        val_loss = self.criterion(y_hat, y)
        self.log('val_loss', val_loss, prog_bar = True)
        return val_loss
        
    def test_step(self, batch, batch_idx):
        x, y = self.prepare_batch(batch)
        y = y.view(y.size(0), -1)
        y = y.type(torch.float)
        y_hat = self.forward(x)
        testAcc = self.accuracy(y_hat, y)
        self.log_dict({'test_acc': testAcc})
        return testAcc

    def configure_optimizers(self):
        optimizer = torch.optim.Adam(self.parameters(), lr=self.learning_rate)
        return optimizer
The odd part is that the validation sanity completes and I get 75% through the first epoch when the validation loop starts before I get the error:

And I have checked that all my inputs have the same dimensions of [193, 229, 193]. Have I missed something obvious? (sorry if it is obvious)
Any help would be greatly appreciated!
For completeness:

pytorch version: 1.9.0
pytorch lightning version: 1.3.7
To me it seems like you have forgotten the batch dimension. 2D convolutions expect input to have shape [N, C, H, W] where C=193, H=229 and W=193 (is it correct that you have the same amount of channels as the width?). If you only want to feed in a single image you can do sample.unsqueeze(0) to add the extra batch dimension in front.