Update diffloss.py

diffloss.py

@@ -91,10 +91,20 @@ class TimestepEmbedder(nn.Module):
             embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1)
         return embedding
 
+    # def forward(self, t):
+    #     t_freq = self.timestep_embedding(t, self.frequency_embedding_size)
+    #     t_emb = self.mlp(t_freq)
+    #     return t_emb
     def forward(self, t):
-        t_freq = self.timestep_embedding(t, self.frequency_embedding_size)
-        t_emb = self.mlp(t_freq)
-        return t_emb
+    t = t.to(self.mlp.weight.device)
+
+    t_freq = self.timestep_embedding(t, self.frequency_embedding_size)
+    
+    t_freq = t_freq.to(self.mlp.weight.device)
+    
+    t_emb = self.mlp(t_freq)
+
+    return t_emb
 
 
 class ResBlock(nn.Module):