fix

main.py

@@ -381,8 +381,10 @@ def train(trainloader, optimizer, model, epoch):
         }
         loss, logits = model(inputs)
         print(loss.shape)
-        loss.backward()
+        print(logits.shape)
         exit()
+        loss.backward()
+        
         optimizer.step()
         pred = logits.argmax(dim=-1)