Merge pull request #397 from salehsargolzaee/my-patch-1

Include model.train() and model.eval()

Author: ronny-udacity

weight-initialization/helpers.py

@@ -28,6 +28,7 @@ def _get_loss_acc(model, train_loader, valid_loader):
         ###################
         # train the model #
         ###################
+        model.train()
         for data, target in train_loader:
             # clear the gradients of all optimized variables
             optimizer.zero_grad()
@@ -45,6 +46,7 @@ def _get_loss_acc(model, train_loader, valid_loader):
     # after training for 2 epochs, check validation accuracy 
     correct = 0
     total = 0
+    model.eval()
     for data, target in valid_loader:
         # forward pass: compute predicted outputs by passing inputs to the model
         output = model(data)
@@ -106,4 +108,4 @@ def hist_dist(title, distribution_tensor, hist_range=(-4, 4)):
     """
     plt.title(title)
     plt.hist(distribution_tensor, np.linspace(*hist_range, num=len(distribution_tensor)//2))
-    plt.show()
+    plt.show()