Thank you everyone (@eqy, @m3tobom_M ). At the moment, my i5 hangs (it is only just a laptop) with this code. I have ironed out the errors but it just hangs and my computer stops. I ran the FashionMNIST in a reasonable time but this is super slow. Can anyone help me out?
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(200*200), transforms.CenterCrop(224), transforms.ToTensor(),]) 
# This gives us some transforms to apply later on

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(200*200, 40000), 
            nn.ReLU(),
            nn.Linear(40000, 10000),
            nn.ReLU(),
            nn.Linear(10000, 75),
            nn.ReLU(),
            nn.Linear(75, 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
check = 0
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
        check = check + 1
        print(check)
        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())

Upon adding some print commands, I get stuck after print(check) and before creating the model. The main one is after the print(check) though. That’s where everything just hangs.