Shape of X [N , C, H, W]: torch.Size([64, 1, 28, 28])
Shape of y: torch.Size([64]) torch.int64
using cuda device
NeuralNetwork(
  (flatten): Flatten(start_dim=1, end_dim=-1)
  (linear_relu_stack): Sequential(
    (0): Linear(in_features=784, out_features=512, bias=True)
    (1): ReLU()
    (2): Linear(in_features=512, out_features=512, bias=True)
    (3): ReLU()
    (4): Linear(in_features=512, out_features=10, bias=True)
  )
)
Epoch 1
------------------------
loss : 2.312647 [   64/60000]
loss : 2.299425 [ 6464/60000]
loss : 2.279648 [12864/60000]
loss : 2.273697 [19264/60000]
loss : 2.257380 [25664/60000]
loss : 2.231007 [32064/60000]
loss : 2.238533 [38464/60000]
loss : 2.205504 [44864/60000]
loss : 2.205637 [51264/60000]
loss : 2.176961 [57664/60000]
Test Error: 
 Accuracy: 41.2%, Avg loss: 2.169525 

Epoch 2
------------------------
loss : 2.180219 [   64/60000]
loss : 2.175783 [ 6464/60000]
loss : 2.121677 [12864/60000]
loss : 2.139049 [19264/60000]
loss : 2.091483 [25664/60000]
loss : 2.028101 [32064/60000]
loss : 2.062609 [38464/60000]
loss : 1.985163 [44864/60000]
loss : 1.990195 [51264/60000]
loss : 1.920861 [57664/60000]
Test Error: 
 Accuracy: 58.5%, Avg loss: 1.919408 

Epoch 3
------------------------
loss : 1.951617 [   64/60000]
loss : 1.927411 [ 6464/60000]
loss : 1.817076 [12864/60000]
loss : 1.854004 [19264/60000]
loss : 1.746010 [25664/60000]
loss : 1.690311 [32064/60000]
loss : 1.714876 [38464/60000]
loss : 1.614491 [44864/60000]
loss : 1.632962 [51264/60000]
loss : 1.525552 [57664/60000]
Test Error: 
 Accuracy: 59.5%, Avg loss: 1.543495 

Epoch 4
------------------------
loss : 1.612580 [   64/60000]
loss : 1.575629 [ 6464/60000]
loss : 1.428641 [12864/60000]
loss : 1.493527 [19264/60000]
loss : 1.373155 [25664/60000]
loss : 1.364947 [32064/60000]
loss : 1.378833 [38464/60000]
loss : 1.300565 [44864/60000]
loss : 1.331147 [51264/60000]
loss : 1.230378 [57664/60000]
Test Error: 
 Accuracy: 62.0%, Avg loss: 1.258319 

Epoch 5
------------------------
loss : 1.339829 [   64/60000]
loss : 1.318503 [ 6464/60000]
loss : 1.155264 [12864/60000]
loss : 1.256479 [19264/60000]
loss : 1.129665 [25664/60000]
loss : 1.156033 [32064/60000]
loss : 1.177031 [38464/60000]
loss : 1.111226 [44864/60000]
loss : 1.145611 [51264/60000]
loss : 1.066547 [57664/60000]
Test Error: 
 Accuracy: 64.1%, Avg loss: 1.087662 

Done!