Unverified Commit 1b2a79d1 authored by Phil Wang's avatar Phil Wang Committed by GitHub
Browse files

Merge pull request #144 from tchamb/import-functools 

Minor patch, couple imports missing
parents 6734f7f8 43def8e0

audiolm_pytorch/encodec.py

+4 −1
Original line number Diff line number Diff line 
from functools import reduce

from einops import rearrange

import torch

from torch import nn

from encodec import EncodecModel

from encodec.utils import convert_audio, _linear_overlap_add
@@ -36,7 +39,7 @@ class EncodecWrapper(nn.Module): 


    @property

    def seq_len_multiple_of(self):

        return functools.reduce(lambda x, y: x * y, self.strides)

        return reduce(lambda x, y: x * y, self.strides)



    def forward(self, x, x_sampling_rate=24000, **kwargs):

        # kwargs for stuff like return_encoded=True, which SoundStream uses but Encodec doesn't

audiolm_pytorch/trainer.py

+1 −1
Original line number Diff line number Diff line
@@ -968,7 +968,7 @@ class CoarseTransformerTrainer(nn.Module): 
        # save model every so often



        if self.is_main and not (steps % self.save_model_every):

            model_path = str(self.results_folder / f'fine.transformer.{steps}.pt')

            model_path = str(self.results_folder / f'coarse.transformer.{steps}.pt')

            self.save(model_path)



            self.print(f'{steps}: saving model to {str(self.results_folder)}')