0

I'm trying to break down the error I'm getting from the loss function. My train function is below:

model = model.to(device)
optimizer = torch.optim.SGD(model.parameters(), lr=LR)
criterion = nn.CrossEntropyLoss()

for epoch in range(2):
    for idx, (img, label) in enumerate(train_dl):
        print(img.shape, label.shape)
        
        img = img.to(device)
        label = label.to(device)
        
        optimizer.zero_grad()
        outputs = model(img)
        
        print(outputs.shape)
        loss = criterion(outputs, label)
        break

The out put is as follows:

img.shape: torch.Size([32, 3, 224, 224]), label.shape torch.Size([32])
outputs.shape: torch.Size([32, 64, 27, 5])

What am I doing wrong? Does it have something to do with my label shape?

zampoan
  • 57
  • 7
  • What are trying to do ? Multiclass classification ? What does you network do, I find it weird the output of your model is (batch size, 64, 27, 5 ) – PlainRavioli Sep 20 '22 at 13:07

0 Answers0