From 1e0b820bdc56a153eac95601c2833a0285957ffc Mon Sep 17 00:00:00 2001 From: SsnL Date: Wed, 31 Jan 2018 13:03:36 -0500 Subject: [PATCH] set transfer learning model in eval mode in visualize_model function --- beginner_source/transfer_learning_tutorial.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/beginner_source/transfer_learning_tutorial.py b/beginner_source/transfer_learning_tutorial.py index e6be2bc678e..de8a5f6fc8e 100644 --- a/beginner_source/transfer_learning_tutorial.py +++ b/beginner_source/transfer_learning_tutorial.py @@ -222,6 +222,8 @@ def train_model(model, criterion, optimizer, scheduler, num_epochs=25): # def visualize_model(model, num_images=6): + was_training = model.training + model.eval() images_so_far = 0 fig = plt.figure() @@ -243,7 +245,9 @@ def visualize_model(model, num_images=6): imshow(inputs.cpu().data[j]) if images_so_far == num_images: + model.train(mode=was_training) return + model.train(mode=was_training) ###################################################################### # Finetuning the convnet