I am trying to train a custom convnet, and I am getting this error:
RuntimeError: mat1 and mat2 shapes cannot be multiplied (1x25088 and 784x16)
From what I have seen, I understand that there is an issue with inputs/outputs with each layer, but I am not sure where exactly is my mistake. Can anyone help me out with figuring out what is causing this issue?
I have included my model class and training loop below.
Model class:
class OsicModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(1,32,kernel_size=(4,4),padding=1,stride=2)
        self.max_pool = nn.MaxPool2d(kernel_size=(2, 2))
        self.conv2 = nn.Conv2d(32, 64, kernel_size=(4, 4), padding=1, stride=2)
        self.bn1 = nn.BatchNorm2d(64)
        self.dropout = nn.Dropout(0.15)
        self.conv3 = nn.Conv2d(64, 4, kernel_size=(4, 4))
        self.bn2 = nn.BatchNorm2d(4)
        self.ada_pool = nn.AdaptiveMaxPool2d((28, 28))
        self.fc1 = nn.Linear(784, 16)
        self.fc2 = nn.Linear(16, 8)
        self.fc3 = nn.Linear(8, 1)
        
        self.relu = nn.ReLU()
        
    def forward(self, x):
        x = self.conv1(x)
        x = self.relu(x)
        x = self.max_pool(x)
        x = self.conv2(x)
        x = self.relu(x)
        x = self.bn1(x)
        x = self.max_pool(x)
        x = self.dropout(x)
        x = self.conv3(x)
        x = self.relu(x)
        x = self.bn2(x)
        x = self.ada_pool(x)
        print(x.shape)
        x = torch.flatten(x)
        x = self.fc1(x)
        x = self.relu(x)
        x = self.fc2(x)
        
        return x

and training loop:
def train(model, train_dataloader, valid_dataloader, test_dataloader, epochs):
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    model.to(device)
    optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
    criterion = nn.MSELoss()
    
    print("start of training loop")
    for epoch in track(range(epochs)):
        model.train()
        train_loss = 0
        
        for batch in train_dataloader:
            optimizer.zero_grad()
            
            imgs = batch['img'].to(device)
            targets = batch['label'].to(device)
            out = model(imgs)
            loss = criterion(out.view(-1, 1), targets)
            loss.backward()
            optimizer.step()
            
            train_loss += loss.item() * targets.size(0)
               
        train_loss /= len(train_dataloader.sampler)

        
        # Validation
        model.eval()
        val_loss = 0

        
        for batch in val_dataloader:
            optimizer.zero_grad()
            
            imgs = batch['img'].to(device)
            targets = batch['target'].to(device)
            
            out = model(imgs)
            loss = criterion(out.view(-1, 1), targets)
            
            val_loss += loss.item() * targets.size(0)
            
        
        val_loss /= len(val_dataloader.sampler)