From 16cc345731954ae575228697b6ce267b67183a35 Mon Sep 17 00:00:00 2001 From: wailovet Date: Fri, 28 Jun 2024 01:44:08 +0800 Subject: [PATCH] 'update' --- hook_HYDiT_main.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/hook_HYDiT_main.py b/hook_HYDiT_main.py index b0d5b6e..d6e3292 100644 --- a/hook_HYDiT_main.py +++ b/hook_HYDiT_main.py @@ -576,11 +576,12 @@ def Core(args, LOG): with torch.cuda.stream(ema_stream): ema.update(model.module.module, step=step) torch.cuda.current_stream().wait_stream(ema_stream) + + with torch.cuda.amp.autocast(enabled=args.use_fp16): + loss_dict = diffusion.training_losses( + model=model, x_start=latents, model_kwargs=model_kwargs) - loss_dict = diffusion.training_losses( - model=model, x_start=latents, model_kwargs=model_kwargs) - - loss = loss_dict["loss"].mean() + loss = loss_dict["loss"].mean() # print(f"step={step}, loss={loss.item()}")