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()
```