diff options
Diffstat (limited to 'trainer')
| -rw-r--r-- | trainer/plugins.py | 5 |
1 files changed, 3 insertions, 2 deletions
diff --git a/trainer/plugins.py b/trainer/plugins.py index 0126870..dc3b24a 100644 --- a/trainer/plugins.py +++ b/trainer/plugins.py @@ -143,18 +143,19 @@ class GeneratorPlugin(Plugin): pattern = 'd-{}-ep{}-s{}.wav' - def __init__(self, samples_path, n_samples, sample_length, sample_rate): + def __init__(self, samples_path, n_samples, sample_length, sample_rate, primer): 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 def register(self, trainer): self.generate = Generator(trainer.model.model, trainer.cuda) def epoch(self, epoch_index): - samples = self.generate(self.n_samples, self.sample_length) \ + samples = self.generate(self.n_samples, self.sample_length, self.primer) \ .cpu().float().numpy() for i in range(self.n_samples): write_wav( |