change in embeddings_flax.py

prishajain1 · prishajain1 · commit f53112d8c795 · 2026-04-12T17:37:38.000+05:30
diff --git a/src/maxdiffusion/models/embeddings_flax.py b/src/maxdiffusion/models/embeddings_flax.py
@@ -521,7 +521,7 @@ def __call__(self, caption, timestep):
     timesteps_proj = self.time_proj(timestep)
     timesteps_emb = self.emb.timestep_embedder(timesteps_proj)
 
-    return hidden_states + timesteps_emb
+    return hidden_states + timesteps_emb[:, None, :]
 
 
 class CombinedTimestepGuidanceTextProjEmbeddings(nn.Module):
