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

Merge pull request #27 from Froskekongen/patch-1 

Update trainer.py
parents 039ec88b 1dc5ab1c

audiolm_pytorch/trainer.py

+1 −1
Original line number Diff line number Diff line
@@ -96,7 +96,7 @@ class SoundStreamTrainer(nn.Module): 
        self.accelerator = Accelerator(**accelerate_kwargs)



        self.soundstream = soundstream

        self.ema_soundstream = EMA(soundstream, update_after_step = ema_update_after_step, update_every = ema_update_every)

        self.ema_soundstream = EMA(soundstream, beta=ema_beta, update_after_step = ema_update_after_step, update_every = ema_update_every)



        self.register_buffer('steps', torch.Tensor([0]))