From 331c40f0c8a47d1a63a313078b70451cb3517f81 Mon Sep 17 00:00:00 2001
From: James Betker <jbetker@gmail.com>
Date: Mon, 19 Oct 2020 15:23:04 -0600
Subject: [PATCH] Allow starting step to be forced

Useful for testing purposes or to force a validation.
---
 codes/train.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/codes/train.py b/codes/train.py
index efc6e378..4da72be3 100644
--- a/codes/train.py
+++ b/codes/train.py
@@ -170,6 +170,8 @@ def main():
     else:
         current_step = -1 if 'start_step' not in opt.keys() else opt['start_step']
         start_epoch = 0
+    if 'force_start_step' in opt.keys():
+        current_step = opt['force_start_step']
 
     #### training
     logger.info('Start training from epoch: {:d}, iter: {:d}'.format(start_epoch, current_step))