fix confidence when alg_temp > 0

generation_utils.py

@@ -441,8 +441,8 @@ class DreamGenerationMixin:
                     if alg_temp is None or alg_temp == 0:
                         _, transfer_index = torch.topk(full_confidence, number_transfer_tokens)
                     else:
-                        confidence = confidence / alg_temp
-                        confidence = F.softmax(confidence, dim=-1)
+                        full_confidence = full_confidence / alg_temp
+                        full_confidence = F.softmax(full_confidence, dim=-1)
                         transfer_index = torch.multinomial(full_confidence, num_samples=number_transfer_tokens)
                     x_ = torch.zeros_like(x, device=self.device, dtype=torch.long) + mask_token_id
                     x_[mask_index] = x0.clone()