From 4f13fb0fee3ca7d2870768e72bc94ed6907b8b8e Mon Sep 17 00:00:00 2001 From: H1yori233 Date: Sun, 28 Jun 2026 11:19:06 +0000 Subject: [PATCH] fix --- fastvideo/train/methods/consistency_model/causal_cd.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/fastvideo/train/methods/consistency_model/causal_cd.py b/fastvideo/train/methods/consistency_model/causal_cd.py index e7a05e5e0..563ad0604 100644 --- a/fastvideo/train/methods/consistency_model/causal_cd.py +++ b/fastvideo/train/methods/consistency_model/causal_cd.py @@ -97,7 +97,7 @@ class CausalConsistencyDistillationMethod(TrainingMethod): sigmas = self._sf_scheduler.sigmas.to(device) timesteps = self._sf_scheduler.timesteps.to(device) - idx = int(torch.randint(0, self._discrete_cd_n - 1, (1, ), generator=self.cuda_generator, device=device).item()) + idx = torch.randint(0, self._discrete_cd_n - 1, (1, ), generator=self.cuda_generator, device=device).squeeze(0) t, t_next = timesteps[idx], timesteps[idx + 1] sigma_t, sigma_t_next = sigmas[idx], sigmas[idx + 1] t_pf = t * torch.ones(batch_size, num_latents, device=device) @@ -149,7 +149,8 @@ class CausalConsistencyDistillationMethod(TrainingMethod): loss = F.mse_loss(x0_t.float(), x0_t_next.float()) loss_map = {"total_loss": loss, "causal_cd_loss": loss} - outputs: dict[str, Any] = {"_fv_backward": (t_pf, training_batch.attn_metadata)} + attn_metadata = (training_batch.attn_metadata_vsa if self._attn_kind == "vsa" else training_batch.attn_metadata) + outputs: dict[str, Any] = {"_fv_backward": (t_pf, attn_metadata)} metrics: dict[str, LogScalar] = {} return loss_map, outputs, metrics