Update diffloss.py

diffloss.py

@@ -96,15 +96,15 @@ class TimestepEmbedder(nn.Module):
     #     t_emb = self.mlp(t_freq)
     #     return t_emb
     def forward(self, t):
-    t = t.to(self.mlp.weight.device)
-
-    t_freq = self.timestep_embedding(t, self.frequency_embedding_size)
+        t = t.to(self.mlp.weight.device)
     
-    t_freq = t_freq.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)
     
-    t_emb = self.mlp(t_freq)
-
-    return t_emb
+        return t_emb
 
 
 class ResBlock(nn.Module):