PyTorch pretrained models not working

Hello, I’m trying to use PyTorch inbuilt models to categorise images, but the loss isn’t decreasing for cross entropy loss and is decreasing for mse loss. I’ve tried adding softmax and different loss functions but only mse loss works. This problem persists for other models like efficient net.
Cross entropy loss

This is the code

weights = AlexNet_Weights.DEFAULT
transform = weights.IMAGENET1K_V1.transforms()
model = alexnet(weights=None)
model.to(device)
data = '/Users/me/Downloads/archive/train'
dataset = datasets.ImageFolder(data, transform=transform)
model.classifier[6] = nn.Linear(4096,450)
dataloader = torch.utils.data.DataLoader(dataset, batch_size=32, shuffle=True, pin_memory=True, num_workers=0)

images, labels = next(iter(dataloader))
loss_function = nn.MSELoss()
optimizer = torch.optim.SGD(model.parameters(),weight_decay=0.005, lr=0.005, momentum=0.9)
model.eval()

for epoch in range(50):
    data_x = []
    data_y = []
    images, labels = next(iter(dataloader))
    for i in range(32):
        data_x.append(np.array(images[i]))
        yes = torch.zeros(450)
        yes[labels[i].item()] = 1
        data_y.append(np.array(yes))
    print(epoch)
    data_x,data_y = np.array(data_x),np.array(data_y)
    data_x,data_y = torch.as_tensor(data_x,device=device),torch.as_tensor(data_y,device=device)
    output_y = model(data_x)
    loss = loss_function(output_y,data_y)
    losses.append(loss.item())
    optimizer.zero_grad(set_to_none=True)
    loss.backward()
    optimizer.step()
    

It seems you are only using the first batch of your DataLoader:

images, labels = next(iter(dataloader))

in each epoch so you might want to change it and then play around with some hyperparameters.

I checked the images and labels by printing them every epoch and they were different, plus I used this same system for previous neural networks and it worked fine