flag to prime the generator

1 files changed, 5 insertions, 1 deletions

diff --git a/train.py b/train.py
index 1f4ab3b..a40e5f6 100644
--- a/train.py
+++ b/train.py
@@ -218,7 +218,8 @@ def main(exp, frame_sizes, dataset, **params):
     ))
     trainer.register_plugin(GeneratorPlugin(
         os.path.join(results_path, 'samples'), params['n_samples'],
-        params['sample_length'], params['sample_rate']
+        params['sample_length'], params['sample_rate'],
+        params['primer']
     ))
     trainer.register_plugin(
         Logger([
@@ -355,6 +356,9 @@ if __name__ == '__main__':
     parser.add_argument(
         '--comet_key', help='comet.ml API key'
     )
+    parser.add_argument(
+        '--primer', help='prime the generator...'
+    )
 
     parser.set_defaults(**default_params)