feat: early return when trained 10 epoch

infer/modules/train/train.py

@@ -248,8 +248,8 @@ def run(rank, n_gpus, hps, logger: logging.Logger, state):
     scaler = GradScaler(enabled=hps.train.fp16_run)
 
     cache = []
+    trained = 0
     for epoch in range(epoch_str, hps.train.epochs + 1):
-        print("epoch", epoch)
         if rank == 0:
             train_and_evaluate(
                 rank,
@@ -283,6 +283,10 @@ def run(rank, n_gpus, hps, logger: logging.Logger, state):
         scheduler_g.step()
         scheduler_d.step()
 
+        trained += 1
+        if trained >= 10:
+            break
+
 
 def train_and_evaluate(
     rank,