Update lm.py

float32 cast during norm

Author: julian-parker

stable_audio_tools/training/lm.py

@@ -236,7 +236,7 @@ def on_train_batch_end(self, trainer, module: AudioLanguageModelTrainingWrapper,
 
                 filename = f'demo_cfg_{cfg_scale}_{trainer.global_step:08}.wav'
                 fakes = fakes / fakes.abs().max()
-                fakes = fakes.clamp(-1, 1).mul(32767).to(torch.int16).cpu()
+                fakes = fakes.type(torch.float32).clamp(-1, 1).mul(32767).type(torch.int16).cpu()
                 torchaudio.save(filename, fakes, self.sample_rate)
 
                 log_dict[f'demo_cfg_{cfg_scale}'] = wandb.Audio(filename,