fix: fix typo & move tensor to cuda device

2023-03-14 02:12:58 +08:00 · 2023-03-14 02:12:58 +08:00 · fab0bc4af1
commit fab0bc4af1
parent dcfeac845e
1 changed files with 2 additions and 2 deletions
--- a/ddpm.py
+++ b/ddpm.py
@ -77,9 +77,9 @@ class DDPM(nn.Module):

                beta = self.beta[time_seq][:, None, None, None].repeat(1, c, h, w)

-                z = torch.randn((n, c, h, w))
+                z = torch.randn((n, c, h, w)).to(self.device)

                x_t = first_term * (x_t-(second_term * predict_noise)) - z * beta
        x_t = ( x_t.clamp(-1, 1) + 1 ) / 2
-        x = x * 255
+        x_t = x_t * 255
        return x_t