Update 06-optim-save-load

Author: yakhyo

Diff for: tutorials/01-basics/06-optim-save-load/main.py

@@ -1,6 +1,6 @@
-# =============================== #
-#           Optimization          #
-# =============================== #
+# ========================================= #
+#               Optimization                #
+# ========================================= #
 
 # Prerequisite Code
 
@@ -65,9 +65,9 @@ def forward(self, x):
 optimizer = torch.optim.SGD(model.parameters(), lr=learning_rate)
 
 
-# =============================== #
-#       Full Implementation       #
-# =============================== #
+# ========================================= #
+#               Full Implementation         #
+# ========================================= #
 
 def train_loop(dataloader, model, loss_fn, optimizer):
     size = len(dataloader.dataset)
@@ -101,19 +101,19 @@ def test_loop(dataloader, model, loss_fn):
     print(f"Test Error: \n Accuracy: {(100 * correct):>0.1f}%, Avg loss: {test_loss:>8f} \n")
 
 
-# =============================== #
-#             Training            #
-# =============================== #
+# ========================================= #
+#                   Training                #
+# ========================================= #
 
 for t in range(epochs):
     print(f"Epoch {t + 1}\n-------------------------------")
     train_loop(train_dataloader, model, loss_fn, optimizer)
     test_loop(test_dataloader, model, loss_fn)
 print("Done!")
 
-# =============================== #
-#     Save and Load the Model     #
-# =============================== #
+# ========================================= #
+#             Save and Load the Model       #
+# ========================================= #
 
 model = models.vgg16(pretrained=True)
 torch.save(model.state_dict(), 'model_weights.pth')  # save the model