Update models/DiT.py

Adjusted TimestepEmbedder.forward to fix tensor dtype inconsistency

models/DiT.py

@@ -22,7 +22,11 @@ class TimestepEmbedder(nn.Module):
         ).to(device=t.device)
         args = torch.einsum('i,j->ij', t, freqs.to(t.device))
         freqs = torch.cat([torch.cos(args),torch.sin(args)],dim=-1)
-        return self.mlp(freqs)
+
+        mlp_dtype = next(self.mlp.parameters()).dtype
+        freqs_casted = freqs.to(mlp_dtype)
+        
+        return self.mlp(freqs_casted)
 
 class ViTAttn(nn.Module):
     def __init__(self,hidden_size,num_heads):