'update'

2024-06-28 01:44:08 +08:00
parent a076b6b4e2
commit 16cc345731
1 changed files with 5 additions and 4 deletions
@@ -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()}")