Fix testloss calculation in Quickstart and Optimization tutorials (pytorch#1536)

* Update optimization_tutorial.py

* Update quickstart_tutorial.py

* Fix typo in optimization tutorial

Co-authored-by: suraj813 <[email protected]>

Author: FaranIdo

beginner_source/basics/optimization_tutorial.py

@@ -167,6 +167,7 @@ def train_loop(dataloader, model, loss_fn, optimizer):
 
 def test_loop(dataloader, model, loss_fn):
     size = len(dataloader.dataset)
+    num_batches = len(dataloader)
     test_loss, correct = 0, 0
 
     with torch.no_grad():
@@ -175,7 +176,7 @@ def test_loop(dataloader, model, loss_fn):
             test_loss += loss_fn(pred, y).item()
             correct += (pred.argmax(1) == y).type(torch.float).sum().item()
 
-    test_loss /= size
+    test_loss /= num_batches
     correct /= size
     print(f"Test Error: \n Accuracy: {(100*correct):>0.1f}%, Avg loss: {test_loss:>8f} \n")
 

beginner_source/basics/quickstart_tutorial.py

@@ -160,6 +160,7 @@ def train(dataloader, model, loss_fn, optimizer):
 
 def test(dataloader, model, loss_fn):
     size = len(dataloader.dataset)
+    num_batches = len(dataloader)
     model.eval()
     test_loss, correct = 0, 0
     with torch.no_grad():
@@ -168,7 +169,7 @@ def test(dataloader, model, loss_fn):
             pred = model(X)
             test_loss += loss_fn(pred, y).item()
             correct += (pred.argmax(1) == y).type(torch.float).sum().item()
-    test_loss /= size
+    test_loss /= num_batches
     correct /= size
     print(f"Test Error: \n Accuracy: {(100*correct):>0.1f}%, Avg loss: {test_loss:>8f} \n")