Update app.py

app.py

@@ -399,14 +399,26 @@ def generate_dream_response(
 
                                 # Normalize probabilities if they don't sum to 1
                                 prob_sum = conf_probs.sum()
-                                if not torch.isclose(prob_sum, torch.tensor(1.0, device=device), atol=1e-4) and prob_sum > 0:
+                                # --- START FIX ---
+                                # Ensure the comparison tensor has the same dtype as prob_sum
+                                target_sum_tensor = torch.tensor(1.0, device=device, dtype=prob_sum.dtype)
+                                if not torch.isclose(prob_sum, target_sum_tensor, atol=1e-4) and prob_sum > 0:
+                                # --- END FIX ---
                                     # print(f"Warning step {i}: Confidence probabilities sum {prob_sum:.4f} != 1. Re-normalizing.")
-                                    conf_probs = conf_probs / prob_sum
-
-                                if conf_probs.numel() > 0 and num_samples > 0 and torch.all(conf_probs >= 0) and torch.isclose(conf_probs.sum(), torch.tensor(1.0, device=device)):
+                                    # Avoid division by zero if prob_sum is extremely small or zero
+                                    safe_prob_sum = torch.max(prob_sum, torch.tensor(1e-12, device=device, dtype=prob_sum.dtype))
+                                    conf_probs = conf_probs / safe_prob_sum # Use safe_prob_sum
+
+                                # Ensure num_samples is valid and probabilities are okay for multinomial
+                                # --- START FIX ---
+                                # Check sum again after potential normalization
+                                final_prob_sum_check = conf_probs.sum()
+                                if conf_probs.numel() > 0 and num_samples > 0 and torch.all(conf_probs >= 0) and torch.isclose(final_prob_sum_check, target_sum_tensor, atol=1e-4):
+                                # --- END FIX ---
                                     try:
                                         transfer_indices_relative = torch.multinomial(conf_probs, num_samples=num_samples, replacement=False)
                                     except RuntimeError as e:
+                                        # [Fallback logic remains the same]
                                         print(f"Warning step {i}: Multinomial sampling failed ('{e}'). Falling back to top-k.")
                                         sort_metric = confidence if alg != 'entropy' else -confidence
                                         k_multinomial_fallback = min(num_samples, sort_metric.numel())
@@ -414,8 +426,9 @@ def generate_dream_response(
                                              _, transfer_indices_relative = torch.topk(sort_metric, k=k_multinomial_fallback)
                                         else:
                                              transfer_indices_relative = torch.tensor([], dtype=torch.long, device=device)
-                                else: # Handle cases where multinomial is not possible
-                                    # print(f"Warning step {i}: Invalid probabilities for multinomial sampling. Falling back to top-k.")
+                                else: # Handle cases where multinomial is not possible (e.g., bad probabilities)
+                                    # [Fallback logic remains the same]
+                                    # print(f"Warning step {i}: Invalid probabilities for multinomial sampling (sum={final_prob_sum_check:.4f}). Falling back to top-k.")
                                     sort_metric = confidence if alg != 'entropy' else -confidence
                                     k_multinomial_fallback = min(num_samples, sort_metric.numel())
                                     if k_multinomial_fallback > 0: