Fix some cosmos fp8 issues.

This commit is contained in:
comfyanonymous
2025-01-16 17:45:37 -05:00
parent cca96a85ae
commit 0aa2368e46
2 changed files with 8 additions and 8 deletions

View File

@@ -293,7 +293,7 @@ class GeneralDIT(nn.Module):
x_B_T_H_W_D = self.x_embedder(x_B_C_T_H_W)
if self.extra_per_block_abs_pos_emb:
extra_pos_emb = self.extra_pos_embedder(x_B_T_H_W_D, fps=fps, device=x_B_C_T_H_W.device)
extra_pos_emb = self.extra_pos_embedder(x_B_T_H_W_D, fps=fps, device=x_B_C_T_H_W.device, dtype=x_B_C_T_H_W.dtype)
else:
extra_pos_emb = None