diff --git a/train_ms.py b/train_ms.py
index f745d70..4a2abc4 100644
--- a/train_ms.py
+++ b/train_ms.py
@@ -1,5 +1,4 @@
 # flake8: noqa: E402
-
 import platform
 import os
 import torch
@@ -582,7 +581,7 @@ def train_and_evaluate(
                         epoch,
                         os.path.join(hps.model_dir, "DUR_{}.pth".format(global_step)),
                     )
-                keep_ckpts = getattr(hps.train, "keep_ckpts", 10)
+                keep_ckpts = config.train_ms_config.keep_ckpts
                 if keep_ckpts > 0:
                     utils.clean_checkpoints(
                         path_to_models=hps.model_dir,