From 4127f7d00aacc7e026858116b78b84187f71e6e4 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Stardust=C2=B7=E5=87=8F?= Date: Tue, 5 Sep 2023 01:02:46 +0800 Subject: [PATCH] Update train_ms.py --- train_ms.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/train_ms.py b/train_ms.py index 958171e..297fea6 100644 --- a/train_ms.py +++ b/train_ms.py @@ -42,7 +42,7 @@ torch.backends.cudnn.allow_tf32 = True # If encontered training problem,please torch.set_float32_matmul_precision('medium') torch.backends.cuda.sdp_kernel("flash") torch.backends.cuda.enable_flash_sdp(True) -torch.backends.cuda.enable_mem_efficient_sdp(True) +torch.backends.cuda.enable_mem_efficient_sdp(True) # Not avaliable if torch version is lower than 2.0 torch.backends.cuda.enable_math_sdp(True) global_step = 0 @@ -165,7 +165,7 @@ def run(rank, n_gpus, hps): epoch_str = max(epoch_str, 1) global_step = (epoch_str - 1) * len(train_loader) - except Exception as e: + except Exception as e: print(e) epoch_str = 1 global_step = 0