It is discuss on many platforms but I am still not getting full understanding of this code.
Below are the inputs.
Number of classes = 5 (0,1,2,3,4)
Input Size  =  torch.Size([8, 3, 512, 512]) ## 8 is a batch-size
**Actual Labels size = torch.Size([8]) **
predicted labels = torch.Size([8, 5])

epoch = 1
from torch.autograd import Variable
criteria = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(params = model.parameters() , lr = .001)
for n in range(epoch):
for index,(images,label) in enumerate(train_loader):
images = Variable(images.to(device))
labels = Variable(labels.to(device))
y_pred = model(images)
print(‘images’,images.shape,‘labels’,labels.shape,‘y_pred’,y_pred.shape)
loss = criteria(labels,y_pred)
optimizer.zero_grad()
loss.backward()
optimizer.step()
#print(images[0])
#print(labels[0])
break;

Please let me know if anyone need more Info.