make training checks params at end of each epoch

PengNi · Nov 23, 2021 · 7733fce · 7733fce
1 parent e9327d8
commit 7733fce
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/deepsignal_plant/train.py b/deepsignal_plant/train.py
@@ -126,7 +126,7 @@ def train(args):
  torch.nn.utils.clip_grad_norm_(model.parameters(), 0.5)
  optimizer.step()
 
- if (i + 1) % args.step_interval == 0:
+ if (i + 1) % args.step_interval == 0 or i == total_step - 1:
  model.eval()
  with torch.no_grad():
  vlosses, vlabels_total, vpredicted_total = [], [], []