Update train.py

k2-fsa · Nov 30, 2024 · b7b6520 · b7b6520
1 parent a1ade8e
commit b7b6520
Showing 1 changed file with 6 additions and 6 deletions.
diff --git a/egs/ljspeech/TTS/vits/train.py b/egs/ljspeech/TTS/vits/train.py
@@ -30,7 +30,7 @@
 from lhotse.cut import Cut
 from lhotse.utils import fix_random_seed
 from tokenizer import Tokenizer
-from torch.cuda.amp import GradScaler, autocast
+from torch.amp import GradScaler, autocast
 from torch.nn.parallel import DistributedDataParallel as DDP
 from torch.optim import Optimizer
 from torch.utils.tensorboard import SummaryWriter
@@ -396,7 +396,7 @@ def save_bad_model(suffix: str = ""):
         loss_info["samples"] = batch_size
 
         try:
-            with autocast(enabled=params.use_fp16):
+            with autocast('cuda',enabled=params.use_fp16):
                 # forward discriminator
                 loss_d, stats_d = model(
                     text=tokens,
@@ -414,7 +414,7 @@ def save_bad_model(suffix: str = ""):
             scaler.scale(loss_d).backward()
             scaler.step(optimizer_d)
 
-            with autocast(enabled=params.use_fp16):
+            with autocast('cuda',enabled=params.use_fp16):
                 # forward generator
                 loss_g, stats_g = model(
                     text=tokens,
@@ -673,7 +673,7 @@ def scan_pessimistic_batches_for_oom(
         )
         try:
             # for discriminator
-            with autocast(enabled=params.use_fp16):
+            with autocast('cuda',enabled=params.use_fp16):
                 loss_d, stats_d = model(
                     text=tokens,
                     text_lengths=tokens_lens,
@@ -686,7 +686,7 @@ def scan_pessimistic_batches_for_oom(
             optimizer_d.zero_grad()
             loss_d.backward()
             # for generator
-            with autocast(enabled=params.use_fp16):
+            with autocast('cuda',enabled=params.use_fp16):
                 loss_g, stats_g = model(
                     text=tokens,
                     text_lengths=tokens_lens,
@@ -838,7 +838,7 @@ def remove_short_and_long_utt(c: Cut):
             params=params,
         )
 
-    scaler = GradScaler(enabled=params.use_fp16, init_scale=1.0)
+    scaler = GradScaler('cuda',enabled=params.use_fp16, init_scale=1.0)
     if checkpoints and "grad_scaler" in checkpoints:
         logging.info("Loading grad scaler state dict")
         scaler.load_state_dict(checkpoints["grad_scaler"])