[Internal/PyT/tacotron2] Convergence update

grzegorz-k-karch · shakandrew · commit c992690e6b3a · 2022-01-27T03:46:29.000-08:00
diff --git a/PyTorch/SpeechSynthesis/Tacotron2/train.py b/PyTorch/SpeechSynthesis/Tacotron2/train.py
@@ -528,7 +528,7 @@ def main():
                                                batch_to_gpu,
                                                args.amp)
 
-        if (epoch % args.epochs_per_checkpoint == 0) and args.bench_class == "":
+        if (epoch % args.epochs_per_checkpoint == 0) and (args.bench_class == "" or args.bench_class == "train"):
             save_checkpoint(model, optimizer, scaler, epoch, model_config,
                             args.output, args.model_name, local_rank, world_size)
         if local_rank == 0:
