fix

main.py

@@ -380,7 +380,9 @@ def train(trainloader, optimizer, model, epoch):
             "label": label
         }
         loss, logits = model(inputs)
+        print(loss.shape)
         loss.backward()
+        exit()
         optimizer.step()
         pred = logits.argmax(dim=-1)