get pl param printing being nice

1 files changed, 5 insertions, 2 deletions

diff --git a/trainer/plugins.py b/trainer/plugins.py
index 562355e..132c33d 100644
--- a/trainer/plugins.py
+++ b/trainer/plugins.py
@@ -143,13 +143,16 @@ class GeneratorPlugin(Plugin):
 
     pattern = 'd-{}-ep{}-s{}.wav'
 
-    def __init__(self, samples_path, n_samples, sample_length, sample_rate, primer):
+    def __init__(self, samples_path, n_samples, sample_length, sample_rate, primer, primer_a, primer_b, recursive):
         super().__init__([(1, 'epoch')])
         self.samples_path = samples_path
         self.n_samples = n_samples
         self.sample_length = sample_length
         self.sample_rate = sample_rate
         self.primer = primer
+        self.primer_a = primer_a
+        self.primer_b = primer_b
+        self.recursive = recursive
 
     def register(self, trainer):
         if self.primer == "":
@@ -158,7 +161,7 @@ class GeneratorPlugin(Plugin):
             self.generate = PrimedGenerator(trainer.model.model, trainer.cuda)
 
     def epoch(self, epoch_index):
-        samples = self.generate(self.n_samples, self.sample_length, self.primer) \
+        samples = self.generate(self.n_samples, self.sample_length, self.primer, self.primer_a, self.primer_b, self.recursive) \
                       .cpu().float().numpy()
         for i in range(self.n_samples):
             write_wav(