Update models/unet.py

models/unet.py

@@ -263,6 +263,7 @@ class TimestepEmbedder(nn.Module):
         self.register_buffer("pe", pe)
 
     def forward(self, x):
+        self.pe = self.pe.cuda()
         return self.pe[x]