Skip to content

Commit 4b894e3

Browse files
pink-redPrathik Rao
authored andcommitted
Fix fine-tuning compatibility with deepspeed (huggingface#816)
1 parent 9f8b13f commit 4b894e3

File tree

1 file changed

+1
-1
lines changed

1 file changed

+1
-1
lines changed

examples/text_to_image/train_text_to_image.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -568,7 +568,7 @@ def collate_fn(examples):
568568

569569
# Predict the noise residual and compute loss
570570
noise_pred = unet(noisy_latents, timesteps, encoder_hidden_states).sample
571-
loss = F.mse_loss(noise_pred, noise, reduction="mean")
571+
loss = F.mse_loss(noise_pred.float(), noise.float(), reduction="mean")
572572

573573
# Gather the losses across all processes for logging (if we use distributed training).
574574
avg_loss = accelerator.gather(loss.repeat(args.train_batch_size)).mean()

0 commit comments

Comments
 (0)