Update train_scratch.py

Author: dragen1860

lesson63-迁移学习-自定义数据集实战/train_scratch.py

@@ -29,6 +29,8 @@
 viz = visdom.Visdom()
 
 def evalute(model, loader):
+    model.eval()
+    
     correct = 0
     total = len(loader.dataset)
 
@@ -58,7 +60,8 @@ def main():
 
             # x: [b, 3, 224, 224], y: [b]
             x, y = x.to(device), y.to(device)
-
+            
+            model.train()
             logits = model(x)
             loss = criteon(logits, y)
 
@@ -94,4 +97,4 @@ def main():
 
 
 if __name__ == '__main__':
-    main()
+    main()