bug fix

2020-06-11 10:48:20 +02:00 · 2020-06-11 10:48:20 +02:00 · b8766c05ce
parent e8a59db839
commit b8766c05ce
1 changed files with 1 additions and 1 deletions
--- a/train.py
+++ b/train.py
@ -237,7 +237,7 @@ def train(model, criterion, optimizer, optimizer_st, scheduler,
                    # save model
                    save_checkpoint(model, optimizer, global_step, epoch, model.decoder.r, OUT_PATH,
                                    optimizer_st=optimizer_st,
-                                    model_loss=loss_dict['postnet_loss'].item())
+                                    model_loss=loss_dict['postnet_loss'])

                # Diagnostic visualizations
                const_spec = postnet_output[0].data.cpu().numpy()