diff --git a/torchgen/_autoheuristic/train_decision.py b/torchgen/_autoheuristic/train_decision.py
index 932baf16e8450..20d4ac4fdc395 100644
--- a/torchgen/_autoheuristic/train_decision.py
+++ b/torchgen/_autoheuristic/train_decision.py
@@ -219,8 +219,8 @@ def eval(name, df):
                             print(
                                 f"new best model with {num_correct} correct and {num_wrong} wrong"
                             )
-                            best_model = model
                             best_model_num_correct = num_correct
+                            best_model = model
                             best_model_safe_proba = safe_proba
                             best_model_unsafe_leaves = unsafe_leaves
 
@@ -403,7 +403,7 @@ def calculate_stats(group):
             if relative_std > 5:
                 times = group["feedback"].tolist()
                 times_str = ", ".join([f"{t:.3f}" for t in sorted(times)])
-                log.debug("High relative std: %f. times=%s", relative_std, times_str)
+                log.error("High relative std: %f. times=%s. group=%s", relative_std, times_str, group.to_dict())
             return pd.Series(
                 {
                     "count": count,
@@ -439,12 +439,15 @@ def get_winner_and_speedup(group):
             winner = sorted_group.iloc[0]["choice"]
             winning_time = sorted_group.iloc[0]["median_execution_time"]
             second_best_time = sorted_group.iloc[1]["median_execution_time"]
-            speedup = second_best_time / winning_time
+            speedup = second_best_time / winning_time if winning_time > 0 else 1.0
             unique_choices = group["choice"].unique()
 
             choice2time = {}
             for row in group.itertuples():
-                choice2time[row.choice] = row.median_execution_time
+                for other_row in group.itertuples():
+                    if row.choice == other_row.choice:
+                        choice2time[row.choice] = row.median_execution_time
+                        break
 
             assert len(unique_choices) == len(group), (
                 f"len(unique_choices) != len(group): {len(unique_choices)} != {len(group)}"
@@ -482,8 +485,9 @@ def add_near_best_configs(df):
                         breakpoint()
                     new_row["actual_winner"] = row["winner"]
                     new_row["winner"] = key
+                    new_rows.append(new_row)
                     if relative_performance >= 0.98:
-                        new_rows.append(new_row)
+                        pass
 
             return pd.DataFrame(new_rows).reset_index(drop=True)