fix device mismatch with loss_scale

2023-11-01 09:29:41 +01:00 · 2023-11-01 09:29:41 +01:00 · c485d4ea60
parent a7343ad190
commit c485d4ea60
1 changed files with 1 additions and 1 deletions
--- a/train.py
+++ b/train.py
@ -963,7 +963,7 @@ def main(args):
                loss_scale = loss_scale * mse_loss_weights

            loss = F.mse_loss(model_pred.float(), target.float(), reduction="none")
-            loss = loss.mean(dim=list(range(1, len(loss.shape)))) * loss_scale
+            loss = loss.mean(dim=list(range(1, len(loss.shape)))) * loss_scale.to(unet.device)
            loss = loss.mean()

            return model_pred, target, loss