update

Author: DN6

src/diffusers/models/embeddings.py

@@ -1594,13 +1594,12 @@ def pool_tokens(x: torch.Tensor, mask: torch.Tensor, *, keepdim=False) -> torch.
         Returns:
             pooled: (B, D) tensor of pooled tokens.
         """
-        input_dtype = x.dtype
         assert x.size(1) == mask.size(1)  # Expected mask to have same length as tokens.
         assert x.size(0) == mask.size(0)  # Expected mask to have same batch size as tokens.
         mask = mask[:, :, None].to(dtype=x.dtype)
         mask = mask / mask.sum(dim=1, keepdim=True).clamp(min=1)
         pooled = (x * mask).sum(dim=1, keepdim=keepdim)
-        return pooled.to(input_dtype)
+        return pooled
 
     def forward(self, x: torch.Tensor, mask: torch.BoolTensor) -> torch.Tensor:
         r"""

src/diffusers/models/transformers/transformer_mochi.py

@@ -136,7 +136,6 @@ def forward(
 
         emb = self.linear(self.silu(emb))
         scale_msa, gate_msa, scale_mlp, gate_mlp = emb.chunk(4, dim=1)
-
         hidden_states = self.norm(hidden_states.to(torch.float32)) * (1 + scale_msa[:, None].to(torch.float32))
         hidden_states = hidden_states.to(hidden_states_dtype)