Also, after doing some work after school,
This is m new code:
import torch
import torch.nn as nn
import torchvision
import torchvision.transforms as transforms
from torch.utils.data import DataLoader

import numpy as np
from PIL import Image
import time

start = time.time()
transform = transforms.Compose([transforms.resize((225,225)),transforms.CenterCrop(224),transforms.ToTensor()]) 
# This gives us some transforms to apply later on
# transforms.Resize(225), 

training = torchvision.datasets.ImageFolder(root = "training_set", transform = transform)
print(training)

train_dataloader = DataLoader(dataset = training, shuffle = True, batch_size = 32)
print(train_dataloader)




print("Before NN creation",start-time.time())
# Creating the model
class NeuralNetwork (nn.Module):
    def __init__(self):
        super(NeuralNetwork, self).__init__()
        self.Linear_Stack = nn.Sequential(
            nn.Linear(50176, 1000), 
            nn.ReLU(),
            nn.Linear(1000, 5),
            nn.ReLU(),
            nn.Linear(5, 1),
            nn.LogSoftmax(dim=1)
        )
    def forward (self, x):
        logits = self.Linear_Stack(x)
        return logits

model = NeuralNetwork()
print(model)

epochs = 1

lr = 0.001
loss = nn.CrossEntropyLoss()
optim = torch.optim.SGD(model.parameters(), lr = lr)
print(loss, optim)  

#creating the main loop
def main(loader, model, loss, optim):
    size = len(loader.dataset)
    print(size)
    for batch_idx, (data, target) in enumerate(loader):
        # Prediction
        
        pred = model.forward(data)
        loss = loss(pred, target)
        #backprop
        optim.zero_grad()
        loss.backward()
        optim.step()
        print(loss)

for x in range(epochs):
    main(train_dataloader, model, loss, optim)
    print(start-time.time())

I get this error
RuntimeError: mat1 and mat2 shapes cannot be multiplied (21504x224 and 50176x1000)

Can anyone point me out on what I’m doing wrong?