fix

main.py

@@ -383,7 +383,8 @@ def train(trainloader, optimizer, model, epoch):
         # print(loss.shape)
         # print(logits.shape)
         # exit()
-        loss.sum().backward()
+        loss = loss.sum()
+        loss.backward()
         
         optimizer.step()
         pred = logits.argmax(dim=-1)