Update Modules/predictor_fsq_flow_convnext.py

Modules/predictor_fsq_flow_convnext.py

@@ -1290,7 +1290,7 @@ class FSQFlowPredictorConvNeXt(nn.Module):
         w = torch.exp(log_w) * dur_mask_1.float() * length_scale
         w = torch.where(dur_mask_1.bool(), torch.clamp(w, min=1.0), w)
         #w = torch.ceil(w).squeeze(1)    # [B, T]  integer durations
-
+        w = w.squeeze(1)    # [B, T]  integer durations
         frame_lengths = w.sum(dim=-1).long()                                        # [B]
         codec_lengths = (frame_lengths + self.codec_stride - 1) // self.codec_stride  # [B]