Update diffloss.py

diffloss.py

@@ -96,13 +96,17 @@ class TimestepEmbedder(nn.Module):
     #     t_emb = self.mlp(t_freq)
     #     return t_emb
     def forward(self, t):
-        t = t.to(self.mlp.weight.device)
+
+        device = next(self.mlp.parameters()).device 
+    
+        t = t.to(device)
     
         t_freq = self.timestep_embedding(t, self.frequency_embedding_size)
-        
-        t_freq = t_freq.to(self.mlp.weight.device)
-        
+    
+        t_freq = t_freq.to(device)
+    
         t_emb = self.mlp(t_freq)
+
     
         return t_emb