Update generation_utils.py

generation_utils.py

@@ -179,20 +179,25 @@ class DreamGenerationMixin:
         num_mask_token = mask_index.sum() / mask_index.shape[0]
         number_transfer_tokens = int(num_mask_token * (1 - s / t)) if step < total_steps - 1 else int(num_mask_token)
 
-        # Update predictions for masked positions only
-        x0 = torch.where(mask_index, x0, x)
-        confidence = torch.where(mask_index, confidence, torch.tensor(-float('inf'), device=x0.device))
+        # Create full confidence tensor matching x dimensions
+        full_confidence = torch.full_like(x, -torch.inf, device=x.device, dtype=confidence.dtype)
+
+        # Create temporary tensor for x0 that matches x dimensions
+        x_temp = torch.zeros_like(x, device=x.device, dtype=torch.long) + mask_token_id
+
+        # Fill masked positions with x0 and confidence
+        x_temp[mask_index] = x0.clone()
+        full_confidence[mask_index] = confidence
 
         # RCR: Select tokens based on cumulative confidence
         for j in range(batch_size):
             if number_transfer_tokens > 0:
-                batch_confidence = confidence[j]
-                batch_mask_index = mask_index[j]
+                batch_full_confidence = full_confidence[j]
 
                 # Select top confident tokens to transfer
-                _, select_indices = torch.topk(batch_confidence, k=number_transfer_tokens, largest=True)
-                x[j, select_indices] = x0[j, select_indices]
-                overtime_confidence[j, select_indices] = batch_confidence[select_indices].clone().float()
+                _, select_indices = torch.topk(batch_full_confidence, k=number_transfer_tokens, largest=True)
+                x[j, select_indices] = x_temp[j, select_indices]
+                overtime_confidence[j, select_indices] = batch_full_confidence[select_indices].clone().float()
 
                 # RCR: Re-mask lowest confidence tokens for next steps
                 if step < total_steps - 1: