Mat1 and mat2 cannot be multiplied

Hello all, I’m trying to follow the tutorial but can’t seem to get a clean run of training the model. The error I get is:
“RuntimeError: mat1 and mat2 shapes cannot be multiplied (49152x256 and 65536x512)”

Now I’ve tried changing the value of the in and out for every available option but nothing will make the mats match up, this includes following the answer to similar topics - such as adding x = x.unsqueeze(0) . Any help would be greatly appreciated, my code is below:

transform = transforms.Compose([
    transforms.Resize((256, 256))
])

dataset_train = CustomImageDataset(annotations_file='train_csv/labels.csv', img_dir='train', transform=transform)
dataset_val = CustomImageDataset(annotations_file='val_csv/labels.csv', img_dir='val', transform=transform)

train_dataloader = DataLoader(dataset_train, batch_size=64, shuffle=True)
val_dataloader = DataLoader(dataset_val, batch_size=64, shuffle=True)


device = (
    "cuda"
    if torch.cuda.is_available()
    else "mps"
    if torch.backends.mps.is_availble()
    else "cpu"
)

class NeuralNetwork(nn.Module):
    def __init__(self):
        super().__init__()
        self.flatten = nn.Flatten()
        self.linear_relu_stack = nn.Sequential(
            nn.Linear(256*256, 512),
            nn.ReLU(),
            nn.Linear(512, 512),
            nn.ReLU(),
            nn.Linear(512, 9),
        )

    def forward(self, x):
        x = x.unsqueeze(0)

        x = x.to(torch.float32) 
        logits = self.linear_relu_stack(x)
        return logits

model = NeuralNetwork().to(device)

Your input needs to have shape of (*, 65536)

So I think the answer here is that I need to represent the image in all three colour spectrums (R G B * 256 * 256) so 3 * 256 * 256 so that when it’s flattened it represents a 2d vector of the entire image across all colour spectrums?

Yes, and that is because you have a linear layer that requires a tensor with 256^2 as the embedding dim.