fix train/distill issue (#215)

Co-authored-by: Peiyuan Zhang <a1286225768@gmail.com>
2025-02-25 08:11:17 -08:00
parent 1e08893546
commit 0be4fc62c9
1 changed files with 3 additions and 0 deletions
@@ -523,6 +523,9 @@ class HYVideoDiffusionTransformer(ModelMixin, ConfigMixin):
    ) -> Union[torch.Tensor, Dict[str, torch.Tensor]]:
        if guidance is None:
            guidance = torch.tensor([6016.0], device=hidden_states.device, dtype=torch.bfloat16)
+        if mask_strategy is None:
+            mask_strategy = [[None] * len(self.heads_num)
+                             for _ in range(len(self.double_blocks) + len(self.single_blocks))]
        img = x = hidden_states
        text_mask = encoder_attention_mask
        t = timestep