flag to prime the generator

1 files changed, 7 insertions, 2 deletions

diff --git a/generate.py b/generate.py
index 92a930f..fb63b09 100644
--- a/generate.py
+++ b/generate.py
@@ -51,7 +51,8 @@ default_params = {
     'sample_length': 80000,
     'loss_smoothing': 0.99,
     'cuda': True,
-    'comet_key': None
+    'comet_key': None,
+    'primer': 'zero'
 }
 
 tag_params = [
@@ -222,7 +223,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(
@@ -359,6 +361,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)