update

examples/rnnoise/step_2_train_model.py

@@ -387,13 +387,13 @@ def main():
                     )
 
                     # save path
-                    epoch_dir = serialization_dir / "epoch-{}".format(epoch_idx)
-                    epoch_dir.mkdir(parents=True, exist_ok=False)
+                    save_dir = serialization_dir / "steps-{}".format(step_idx)
+                    save_dir.mkdir(parents=True, exist_ok=False)
 
                     # save models
-                    model.save_pretrained(epoch_dir.as_posix())
+                    model.save_pretrained(save_dir.as_posix())
 
-                    model_list.append(epoch_dir)
+                    model_list.append(save_dir)
                     if len(model_list) >= args.num_serialized_models_to_keep:
                         model_to_delete: Path = model_list.pop(0)
                         shutil.rmtree(model_to_delete.as_posix())
@@ -418,7 +418,7 @@ def main():
                         "pesq_score": average_pesq_score,
                         "loss": average_loss,
                     }
-                    metrics_filename = epoch_dir / "metrics_epoch.json"
+                    metrics_filename = save_dir / "metrics_epoch.json"
                     with open(metrics_filename, "w", encoding="utf-8") as f:
                         json.dump(metrics, f, indent=4, ensure_ascii=False)
 
@@ -427,7 +427,7 @@ def main():
                     if best_epoch_idx == epoch_idx:
                         if best_dir.exists():
                             shutil.rmtree(best_dir)
-                        shutil.copytree(epoch_dir, best_dir)
+                        shutil.copytree(save_dir, best_dir)
 
                     # early stop
                     early_stop_flag = False
@@ -446,5 +446,5 @@ def main():
     return
 
 
-if __name__ == '__main__':
+if __name__ == "__main__":
     main()