update curve

layumi · May 11, 2023 · a64c5b6 · a64c5b6
1 parent 879baf5
commit a64c5b6
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/train.py b/train.py
@@ -397,7 +397,8 @@ def train_model(model, criterion, optimizer, scheduler, num_epochs=25):
                     save_network(model.module, epoch+1)
                 else:
                     save_network(model, epoch+1)
-            draw_curve(epoch)
+            if phase == 'val':
+                draw_curve(epoch)
             if phase == 'train':
                scheduler.step()
         time_elapsed = time.time() - since